"""Tests for U6: attachment & image field upload, download, and cleanup. Requires PostgreSQL — marked ``postgres``. Uses ``httpx.AsyncClient`` with ``ASGITransport`` (same pattern as test_routes.py). """ from __future__ import annotations import io from pathlib import Path from typing import Any import httpx import pytest from fastapi import FastAPI from httpx import ASGITransport from agentkit.bitable.service import BitableService from agentkit.server.routes import bitable as bitable_routes from agentkit.server.routes.bitable import require_bitable_auth pytestmark = pytest.mark.postgres TEST_USER_ID = "test-user-id" def _make_test_user() -> dict[str, Any]: return {"user_id": TEST_USER_ID, "username": "testuser", "role": "member"} @pytest.fixture def app( bitable_service: BitableService, tmp_path: Path, monkeypatch: pytest.MonkeyPatch ) -> FastAPI: """Test app with upload dir redirected to tmp_path.""" upload_dir = tmp_path / "bitable_uploads" # Patch both the routes module variable AND the env var (service reads env var) monkeypatch.setattr(bitable_routes, "BITABLE_UPLOAD_DIR", upload_dir) monkeypatch.setenv("AGENTKIT_BITABLE_UPLOAD_DIR", str(upload_dir)) app = FastAPI() app.state.bitable_service = bitable_service app.include_router(bitable_routes.router, prefix="/api/v1") app.dependency_overrides[require_bitable_auth] = lambda: _make_test_user() return app @pytest.fixture async def client(app: FastAPI) -> httpx.AsyncClient: transport = ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: yield c # --------------------------------------------------------------------------- # Helpers # --------------------------------------------------------------------------- async def _create_table_with_field( client: httpx.AsyncClient, field_type: str, field_name: str = "files", ) -> tuple[str, str]: """Create a table + a field, return (table_id, field_id).""" table_id = (await client.post("/api/v1/bitable/tables", json={"name": "T"})).json()["table"][ "id" ] field_id = ( await client.post( f"/api/v1/bitable/tables/{table_id}/fields", json={"name": field_name, "field_type": field_type, "owner": "agent"}, ) ).json()["field"]["id"] return table_id, field_id def _make_image_bytes(name: str = "test.png", size: int = 100) -> tuple[bytes, str]: """Minimal valid PNG header + padding.""" png_header = b"\x89PNG\r\n\x1a\n" body = b"\x00" * size return png_header + body, name def _make_pdf_bytes(name: str = "doc.pdf", size: int = 50) -> tuple[bytes, str]: return b"%PDF-1.4\n" + b"\x00" * size, name # --------------------------------------------------------------------------- # Upload tests # --------------------------------------------------------------------------- async def test_upload_image_success(client: httpx.AsyncClient, tmp_path: Path) -> None: table_id, field_id = await _create_table_with_field(client, "image") img_bytes, img_name = _make_image_bytes() resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": (img_name, io.BytesIO(img_bytes), "image/png")}, ) assert resp.status_code == 200 data = resp.json() assert data["filename"] == img_name assert data["mime_type"] == "image/png" assert data["size"] == len(img_bytes) assert data["stored_name"].endswith(".png") assert data["url"].startswith("/api/v1/bitable/files/") # File exists on disk file_path = bitable_routes.BITABLE_UPLOAD_DIR / data["stored_name"] assert file_path.exists() assert file_path.read_bytes() == img_bytes async def test_upload_attachment_pdf(client: httpx.AsyncClient) -> None: table_id, field_id = await _create_table_with_field(client, "attachment") pdf_bytes, pdf_name = _make_pdf_bytes() resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": (pdf_name, io.BytesIO(pdf_bytes), "application/pdf")}, ) assert resp.status_code == 200 data = resp.json() assert data["filename"] == pdf_name assert data["mime_type"] == "application/pdf" async def test_upload_image_rejects_non_image(client: httpx.AsyncClient) -> None: table_id, field_id = await _create_table_with_field(client, "image") pdf_bytes, _ = _make_pdf_bytes() resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": ("doc.pdf", io.BytesIO(pdf_bytes), "application/pdf")}, ) assert resp.status_code == 400 assert "image" in resp.json()["detail"].lower() async def test_upload_rejects_non_attachment_field(client: httpx.AsyncClient) -> None: table_id, field_id = await _create_table_with_field(client, "text") img_bytes, _ = _make_image_bytes() resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": ("test.png", io.BytesIO(img_bytes), "image/png")}, ) assert resp.status_code == 400 async def test_upload_404_unknown_field(client: httpx.AsyncClient) -> None: table_id = (await client.post("/api/v1/bitable/tables", json={"name": "T"})).json()["table"][ "id" ] img_bytes, _ = _make_image_bytes() resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": "nonexistent"}, files={"file": ("test.png", io.BytesIO(img_bytes), "image/png")}, ) assert resp.status_code == 404 async def test_upload_requires_auth(bitable_service: BitableService) -> None: """No auth override → 401.""" app = FastAPI() app.state.bitable_service = bitable_service app.include_router(bitable_routes.router, prefix="/api/v1") transport = ASGITransport(app=app) async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: resp = await c.post( "/api/v1/bitable/tables/x/upload", params={"field_id": "y"}, files={"file": ("t.png", io.BytesIO(b"x"), "image/png")}, ) assert resp.status_code == 401 # --------------------------------------------------------------------------- # Download tests # --------------------------------------------------------------------------- async def test_download_file_success(client: httpx.AsyncClient) -> None: table_id, field_id = await _create_table_with_field(client, "image") img_bytes, img_name = _make_image_bytes() upload_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": (img_name, io.BytesIO(img_bytes), "image/png")}, ) stored_name = upload_resp.json()["stored_name"] resp = await client.get(f"/api/v1/bitable/files/{stored_name}") assert resp.status_code == 200 assert resp.content == img_bytes async def test_download_404_missing_file(client: httpx.AsyncClient) -> None: resp = await client.get("/api/v1/bitable/files/nonexistent.png") assert resp.status_code == 404 # --------------------------------------------------------------------------- # Attachment cleanup on record deletion # --------------------------------------------------------------------------- async def test_delete_record_cleans_up_files( client: httpx.AsyncClient, bitable_service: BitableService, ) -> None: table_id, field_id = await _create_table_with_field(client, "image") img_bytes, _ = _make_image_bytes() upload_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": ("pic.png", io.BytesIO(img_bytes), "image/png")}, ) file_meta = upload_resp.json() stored_name = file_meta["stored_name"] # Create a record with the image metadata create_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/records", json={"records": [{field_id: [file_meta]}]}, ) record_id = create_resp.json()["records"][0]["id"] # Verify file exists file_path = bitable_routes.BITABLE_UPLOAD_DIR / stored_name assert file_path.exists() # Delete the record del_resp = await client.delete(f"/api/v1/bitable/records/{record_id}") assert del_resp.status_code == 200 # File should be gone assert not file_path.exists() async def test_delete_records_by_table_cleans_up_files( client: httpx.AsyncClient, ) -> None: table_id, field_id = await _create_table_with_field(client, "attachment") pdf_bytes, _ = _make_pdf_bytes() upload_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": ("doc.pdf", io.BytesIO(pdf_bytes), "application/pdf")}, ) file_meta = upload_resp.json() stored_name = file_meta["stored_name"] await client.post( f"/api/v1/bitable/tables/{table_id}/records", json={"records": [{field_id: [file_meta]}]}, ) file_path = bitable_routes.BITABLE_UPLOAD_DIR / stored_name assert file_path.exists() # Delete all records resp = await client.delete(f"/api/v1/bitable/tables/{table_id}/records") assert resp.status_code == 200 assert not file_path.exists() async def test_delete_record_when_file_already_missing( client: httpx.AsyncClient, ) -> None: """Record deletion should succeed even if the physical file is gone.""" table_id, field_id = await _create_table_with_field(client, "image") img_bytes, _ = _make_image_bytes() upload_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": ("pic.png", io.BytesIO(img_bytes), "image/png")}, ) file_meta = upload_resp.json() stored_name = file_meta["stored_name"] create_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/records", json={"records": [{field_id: [file_meta]}]}, ) record_id = create_resp.json()["records"][0]["id"] # Manually delete the file before deleting the record file_path = bitable_routes.BITABLE_UPLOAD_DIR / stored_name file_path.unlink() assert not file_path.exists() # Record deletion should still succeed del_resp = await client.delete(f"/api/v1/bitable/records/{record_id}") assert del_resp.status_code == 200 # --------------------------------------------------------------------------- # Multiple files in one field # --------------------------------------------------------------------------- async def test_multiple_files_in_attachment_field(client: httpx.AsyncClient) -> None: table_id, field_id = await _create_table_with_field(client, "attachment") metas = [] for name in ("a.pdf", "b.pdf"): pdf_bytes, _ = _make_pdf_bytes(name) resp = await client.post( f"/api/v1/bitable/tables/{table_id}/upload", params={"field_id": field_id}, files={"file": (name, io.BytesIO(pdf_bytes), "application/pdf")}, ) metas.append(resp.json()) # Store all files as an array in one record create_resp = await client.post( f"/api/v1/bitable/tables/{table_id}/records", json={"records": [{field_id: metas}]}, ) assert create_resp.status_code == 201 record = create_resp.json()["records"][0] assert len(record["values"][field_id]) == 2