123 lines
4.4 KiB
Python
123 lines
4.4 KiB
Python
import io
|
|
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from typing import Any
|
|
from unittest.mock import patch
|
|
from uuid import UUID
|
|
|
|
import laspy
|
|
import numpy as np
|
|
|
|
from B03_FileInput.B03_FileInput_Router import upload_project_files
|
|
from B03_FileInput.B03_FileInput_Schema import FileUploadResponse
|
|
|
|
|
|
class UploadProjectFilesTest(unittest.IsolatedAsyncioTestCase):
|
|
async def test_upload_project_files(self) -> None:
|
|
class AsyncContext:
|
|
def __init__(self, value: Any = None) -> None:
|
|
self.value = value
|
|
|
|
async def __aenter__(self) -> Any:
|
|
return self.value
|
|
|
|
async def __aexit__(self, *_args: Any) -> None:
|
|
return None
|
|
|
|
class FakeCursor:
|
|
def __init__(self, connection: "FakeConnection") -> None:
|
|
self.connection = connection
|
|
self.lastrowid = 0
|
|
|
|
async def __aenter__(self) -> "FakeCursor":
|
|
return self
|
|
|
|
async def __aexit__(self, *_args: Any) -> None:
|
|
return None
|
|
|
|
async def execute(self, query: str, _arguments: Any = None) -> None:
|
|
if query.strip().upper().startswith("SELECT"):
|
|
self._row: tuple[Any, ...] | None = ("storage/company/user/project-id",)
|
|
else:
|
|
self.lastrowid = self.connection.next_id
|
|
self.connection.next_id += 1
|
|
self._row = None
|
|
|
|
async def fetchone(self) -> tuple[Any, ...] | None:
|
|
return getattr(self, "_row", None)
|
|
|
|
class FakeConnection:
|
|
def __init__(self) -> None:
|
|
self.next_id = 1
|
|
self.committed = False
|
|
self.rolled_back = False
|
|
|
|
def cursor(self) -> FakeCursor:
|
|
return FakeCursor(self)
|
|
|
|
async def begin(self) -> None:
|
|
return None
|
|
|
|
async def commit(self) -> None:
|
|
self.committed = True
|
|
|
|
async def rollback(self) -> None:
|
|
self.rolled_back = True
|
|
|
|
class FakePool:
|
|
def __init__(self, connection: FakeConnection) -> None:
|
|
self.connection = connection
|
|
|
|
def acquire(self) -> AsyncContext:
|
|
return AsyncContext(self.connection)
|
|
|
|
class FakeUpload:
|
|
def __init__(self, filename: str, content: bytes) -> None:
|
|
self.filename = filename
|
|
self.size = len(content)
|
|
self.file = io.BytesIO(content)
|
|
self.closed = False
|
|
|
|
async def read(self, size: int) -> bytes:
|
|
return self.file.read(size)
|
|
|
|
async def close(self) -> None:
|
|
self.file.close()
|
|
self.closed = True
|
|
|
|
with tempfile.TemporaryDirectory() as temporary_directory:
|
|
project_root = Path(temporary_directory) / "project"
|
|
project_root.mkdir()
|
|
source = Path(temporary_directory) / "source.las"
|
|
las = laspy.LasData(laspy.LasHeader(point_format=3, version="1.2"))
|
|
las.x = np.array([1.0])
|
|
las.y = np.array([2.0])
|
|
las.z = np.array([3.0])
|
|
las.write(source)
|
|
upload = FakeUpload("cloud.las", source.read_bytes())
|
|
pool = FakePool(FakeConnection())
|
|
|
|
with (
|
|
patch("B03_FileInput.B03_FileInput_Router.get_db_pool", return_value=pool),
|
|
patch(
|
|
"B03_FileInput.B03_FileInput_Router.resolve_stored_project_path",
|
|
return_value=str(project_root),
|
|
),
|
|
):
|
|
response = await upload_project_files( # type: ignore[arg-type]
|
|
UUID("550e8400-e29b-41d4-a716-446655440000"),
|
|
[upload],
|
|
)
|
|
|
|
self.assertIsInstance(response, FileUploadResponse)
|
|
assert isinstance(response, FileUploadResponse)
|
|
self.assertEqual(len(response.files), 1)
|
|
self.assertEqual(response.files[0].input_file_id, 1)
|
|
self.assertTrue(
|
|
(project_root / "B03_FileInput" / "input" / "las" / "cloud.las").is_file()
|
|
)
|
|
self.assertTrue((project_root / "B03_FileInput" / "metadata.json").is_file())
|
|
self.assertTrue((project_root / "workflow.json").is_file())
|
|
self.assertTrue(upload.closed)
|