import gc import unittest import warnings from pathlib import Path from tempfile import TemporaryDirectory from unittest.mock import patch 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_non_file_form_fields(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", files={"file": ("quote.pdf", b"pdf", "application/pdf")}, data={"description": "quote"}, ) self.assertEqual(response.status_code, 400) self.assertEqual(response.json()["error"]["code"], "invalid_upload") self.assertEqual(http_client.calls, []) self.assertEqual(list(Path(temp_dir).iterdir()), []) async def test_upload_rejects_oversized_misreported_chunked_body_before_multipart_parsing(self) -> None: http_client = FakeHttpClient() endpoint = UploadEndpoint(http_client, max_bytes=5, temp_dir="/tmp") parsed = False async def forbidden_form(self): nonlocal parsed parsed = True self.fail("multipart parsing must not begin for an oversized body") body = b"x" * (5 + 65537) scope = { "type": "http", "method": "POST", "path": "/uploads", "headers": [ (b"content-type", b"multipart/form-data; boundary=test"), (b"content-length", b"1"), ], } messages = [ {"type": "http.request", "body": body[:10], "more_body": True}, {"type": "http.request", "body": body[10:], "more_body": False}, ] async def receive(): return messages.pop(0) from starlette.requests import Request starlette_request = Request(scope, receive=receive) with patch("nocobase_crm_mcp.uploads.Request.form", new=forbidden_form): response = await endpoint(starlette_request) self.assertEqual(response.status_code, 413) self.assertFalse(parsed) self.assertEqual(http_client.calls, []) 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 = ( (NocoBaseHttpError(401, "upstream failure"), "nb_upload_unauthorized"), (NocoBaseHttpError(403, "expired", authentication_failed=True), "nb_upload_unauthorized"), (NocoBaseHttpError(500, "upstream failure"), "nb_upload_failed"), ) for upstream_error, error_code in cases: with ( self.subTest(status_code=upstream_error.status_code, error_code=error_code), TemporaryDirectory() as temp_dir, ): http_client = FakeHttpClient() http_client.error = upstream_error 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_maps_project_write_denial_to_forbidden(self) -> None: http_client = FakeHttpClient() http_client.error = NocoBaseHttpError(403, "permission denied") 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, 403) self.assertEqual(response.json()["error"]["code"], "nb_upload_forbidden") 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()