215 lines
8.9 KiB
Python
215 lines
8.9 KiB
Python
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, "permission denied"), "nb_upload_failed"),
|
|
(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_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()
|