323 lines
11 KiB
Python
323 lines
11 KiB
Python
"""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 == 200
|
|
record = create_resp.json()["records"][0]
|
|
assert len(record["values"][field_id]) == 2
|