nocobase-mcp/tests/test_uploads.py

import gc
import unittest
import warnings
from pathlib import Path
from tempfile import TemporaryDirectory

import httpx
from starlette.applications import Starlette
from starlette.routing import Route

from nocobase_crm_mcp.nocobase_http import NocoBaseHttpError
from nocobase_crm_mcp.uploads import UploadEndpoint


AUTH = {"Authorization": "Bearer secret"}


class FakeHttpClient:
    def __init__(self) -> None:
        self.calls: list[tuple[Path, str, str | None, bytes]] = []
        self.error: Exception | None = None

    def upload_attachment(self, path: Path, filename: str, mimetype: str | None) -> dict[str, object]:
        self.calls.append((path, filename, mimetype, path.read_bytes()))
        if self.error:
            raise self.error
        return {"id": 42}


class UploadEndpointTests(unittest.IsolatedAsyncioTestCase):
    async def _client_for(
        self,
        http_client: FakeHttpClient,
        temp_dir: str,
        *,
        max_bytes: int = 5,
        allowed_mime_types: tuple[str, ...] | None = None,
    ) -> httpx.AsyncClient:
        endpoint = UploadEndpoint(http_client, max_bytes, temp_dir, allowed_mime_types)

        async def upload_route(request):
            return await endpoint(request)

        app = Starlette(
            routes=[
                Route(
                    "/uploads",
                    upload_route,
                    methods=["GET", "POST"],
                )
            ]
        )
        return httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://testserver")

    async def test_upload_returns_nocobase_attachment_id_and_removes_temp_file(self) -> None:
        http_client = FakeHttpClient()
        with TemporaryDirectory() as temp_dir:
            async with await self._client_for(http_client, temp_dir) as client:
                response = await client.post(
                    "/uploads",
                    headers=AUTH,
                    files={"file": ("quote.pdf", b"pdf", "application/pdf")},
                )

            self.assertEqual(response.status_code, 200)
            self.assertEqual(
                response.json(),
                {"upload_id": 42, "filename": "quote.pdf", "size": 3, "mimetype": "application/pdf"},
            )
            self.assertEqual(http_client.calls[0][1:], ("quote.pdf", "application/pdf", b"pdf"))
            self.assertEqual(list(Path(temp_dir).iterdir()), [])

    async def test_upload_rejects_missing_multiple_and_over_limit_files(self) -> None:
        http_client = FakeHttpClient()
        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always", ResourceWarning)
            with TemporaryDirectory() as temp_dir:
                async with await self._client_for(http_client, temp_dir) as client:
                    missing = await client.post("/uploads", headers=AUTH, files={})
                    multiple = await client.post(
                        "/uploads",
                        headers=AUTH,
                        files=[
                            ("file", ("one.pdf", b"one", "application/pdf")),
                            ("file", ("two.pdf", b"two", "application/pdf")),
                        ],
                    )
                    over_limit = await client.post(
                        "/uploads",
                        headers=AUTH,
                        files={"file": ("big.bin", b"x" * 6, "application/octet-stream")},
                    )
                temp_paths = list(Path(temp_dir).iterdir())
            gc.collect()

        self.assertEqual(missing.status_code, 400)
        self.assertEqual(missing.json()["error"]["code"], "invalid_upload")
        self.assertEqual(multiple.status_code, 400)
        self.assertEqual(multiple.json()["error"]["code"], "invalid_upload")
        self.assertEqual(over_limit.status_code, 413)
        self.assertEqual(over_limit.json()["error"]["code"], "upload_too_large")
        self.assertEqual(http_client.calls, [])
        self.assertEqual(temp_paths, [])
        self.assertFalse(any(issubclass(warning.category, ResourceWarning) for warning in caught))

    async def test_upload_rejects_disallowed_mime_type_before_calling_nocobase(self) -> None:
        http_client = FakeHttpClient()
        with TemporaryDirectory() as temp_dir:
            async with await self._client_for(
                http_client, temp_dir, allowed_mime_types=("image/*", "application/pdf")
            ) as client:
                response = await client.post(
                    "/uploads",
                    headers=AUTH,
                    files={"file": ("notes.txt", b"notes", "text/plain")},
                )

            self.assertEqual(response.status_code, 415)
            self.assertEqual(response.json()["error"]["code"], "unsupported_media_type")
            self.assertEqual(http_client.calls, [])
            self.assertEqual(list(Path(temp_dir).iterdir()), [])

    async def test_upload_maps_upstream_errors_and_removes_temp_file(self) -> None:
        cases = ((401, "nb_upload_unauthorized"), (403, "nb_upload_unauthorized"), (500, "nb_upload_failed"))

        for status_code, error_code in cases:
            with self.subTest(status_code=status_code), TemporaryDirectory() as temp_dir:
                http_client = FakeHttpClient()
                http_client.error = NocoBaseHttpError(status_code, "upstream failure")
                async with await self._client_for(http_client, temp_dir) as client:
                    response = await client.post(
                        "/uploads",
                        headers=AUTH,
                        files={"file": ("quote.pdf", b"pdf", "application/pdf")},
                    )

                self.assertEqual(response.status_code, 502)
                self.assertEqual(response.json()["error"]["code"], error_code)
                self.assertEqual(list(Path(temp_dir).iterdir()), [])

    async def test_upload_requires_post_and_multipart_form_data(self) -> None:
        http_client = FakeHttpClient()
        with TemporaryDirectory() as temp_dir:
            async with await self._client_for(http_client, temp_dir) as client:
                get_response = await client.get("/uploads", headers=AUTH)
                json_response = await client.post("/uploads", headers=AUTH, json={"file": "not-a-file"})

            self.assertEqual(get_response.status_code, 405)
            self.assertEqual(json_response.status_code, 400)
            self.assertEqual(json_response.json()["error"]["code"], "invalid_upload")
            self.assertEqual(list(Path(temp_dir).iterdir()), [])


if __name__ == "__main__":
    unittest.main()