theo-agent-dashboard

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README

test_attachments.py (12609B)


      1 """T059 — Tests for attachment upload and download endpoints."""
      2 import base64
      3 import struct
      4 import zlib
      5 from pathlib import Path
      6 from unittest.mock import patch
      7 
      8 import pytest
      9 from httpx import ASGITransport, AsyncClient
     10 
     11 from app.main import app
     12 
     13 # ---------------------------------------------------------------------------
     14 # Helpers
     15 # ---------------------------------------------------------------------------
     16 def _make_image_png_bytes() -> bytes:
     17     """Return a minimal valid 1×1 white PNG."""
     18     def _chunk(ctype: bytes, data: bytes) -> bytes:
     19         raw = ctype + data
     20         return (
     21             struct.pack(">I", len(data))
     22             + raw
     23             + struct.pack(">I", zlib.crc32(raw) & 0xFFFFFFFF)
     24         )
     25 
     26     ihdr = struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0)
     27     raw_scanline = b"\x00\xff\xff\xff"
     28     idat = zlib.compress(raw_scanline)
     29     return (
     30         b"\x89PNG\r\n\x1a\n"
     31         + _chunk(b"IHDR", ihdr)
     32         + _chunk(b"IDAT", idat)
     33         + _chunk(b"IEND", b"")
     34     )
     35 
     36 
     37 def _make_large_file(size_mb: int) -> bytes:
     38     """Return a file of approximately ``size_mb`` megabytes."""
     39     return b"\x00" * (size_mb * 1024 * 1024 + 1)
     40 
     41 
     42 # ---------------------------------------------------------------------------
     43 # Fixtures
     44 # ---------------------------------------------------------------------------
     45 @pytest.fixture(autouse=True)
     46 def _isolate_env(tmp_path):
     47     """Isolate config, db, and auth for each test."""
     48     from app import config as cfg_mod
     49     from app.db import Database
     50     import app.routes as routes_mod
     51 
     52     # Create temp data dirs
     53     data_dir = tmp_path / "data"
     54     data_dir.mkdir(parents=True)
     55     (data_dir / "attachments").mkdir()
     56 
     57     # Patch config
     58     old_cfg = cfg_mod._config
     59     cfg_mod._config = cfg_mod.DashboardConfig(data_dir=str(data_dir))
     60 
     61     # Create temp database with the attachments table
     62     db = Database(str(tmp_path / "test.db"))
     63     db.connection = db._create_connection()
     64     db.connection.executescript("""
     65         CREATE TABLE IF NOT EXISTS attachments (
     66             id TEXT PRIMARY KEY,
     67             session_id TEXT NOT NULL,
     68             filename TEXT NOT NULL,
     69             content_type TEXT NOT NULL,
     70             size_bytes INTEGER NOT NULL,
     71             disk_path TEXT NOT NULL,
     72             created_at INTEGER NOT NULL
     73         );
     74     """)
     75 
     76     # Patch db in the attachment route
     77     import app.routes.attachments as att_mod
     78     old_route_db = getattr(att_mod, "_db", None)
     79     att_mod._db = db
     80 
     81     # Bypass auth for testing attachment endpoints
     82     old_verify = routes_mod.verify_jwt
     83     old_get_jwt = routes_mod.get_jwt_from_cookie
     84     routes_mod.get_jwt_from_cookie = lambda request: "fake-token"
     85     routes_mod.verify_jwt = lambda token: {"sub": "admin"}
     86 
     87     yield
     88 
     89     # Restore everything
     90     att_mod._db = old_route_db
     91     routes_mod.verify_jwt = old_verify
     92     routes_mod.get_jwt_from_cookie = old_get_jwt
     93     if db.connection:
     94         db.connection.close()
     95     cfg_mod._config = old_cfg
     96 
     97 
     98 @pytest.fixture
     99 async def client() -> AsyncClient:
    100     """Async HTTP client for testing."""
    101     async with AsyncClient(
    102         transport=ASGITransport(app=app),
    103         base_url="http://test",
    104     ) as c:
    105         yield c
    106 
    107 
    108 # ---------------------------------------------------------------------------
    109 # Tests — Upload
    110 # ---------------------------------------------------------------------------
    111 
    112 class TestUploadAttachment:
    113     """POST /api/attachments — file upload."""
    114 
    115     async def test_upload_png_returns_attachment_response(
    116         self, client: AsyncClient,
    117     ) -> None:
    118         """Upload a PNG and verify AttachmentResponse fields."""
    119         png_data = _make_image_png_bytes()
    120         resp = await client.post(
    121             "/api/attachments",
    122             files={"file": ("photo.png", png_data, "image/png")},
    123             data={"session_id": "sess-1"},
    124         )
    125         assert resp.status_code == 200
    126         body = resp.json()
    127         assert body["id"]
    128         assert isinstance(body["id"], str)
    129         assert len(body["id"]) > 0
    130         assert body["filename"] == "photo.png"
    131         assert body["content_type"] == "image/png"
    132         assert body["size_bytes"] == len(png_data)
    133         assert isinstance(body["created_at"], int)
    134         assert body["created_at"] > 0
    135 
    136     async def test_upload_returns_413_for_over_50mb(
    137         self, client: AsyncClient,
    138     ) -> None:
    139         """Files exceeding 50 MB are rejected with 413."""
    140         large_data = _make_large_file(51)
    141         resp = await client.post(
    142             "/api/attachments",
    143             files={"file": ("huge.bin", large_data, "application/octet-stream")},
    144             data={"session_id": "sess-1"},
    145         )
    146         assert resp.status_code == 413
    147         body = resp.json()
    148         assert "detail" in body
    149         assert "50 MB" in body["detail"]
    150 
    151     async def test_upload_txt_file_returns_metadata(
    152         self, client: AsyncClient,
    153     ) -> None:
    154         """Non-image files store on disk and return metadata."""
    155         content = b"hello world"
    156         resp = await client.post(
    157             "/api/attachments",
    158             files={"file": ("note.txt", content, "text/plain")},
    159             data={"session_id": "sess-1"},
    160         )
    161         assert resp.status_code == 200
    162         body = resp.json()
    163         assert body["filename"] == "note.txt"
    164         assert body["content_type"] == "text/plain"
    165         assert body["size_bytes"] == len(content)
    166         assert isinstance(body["id"], str)
    167         assert len(body["id"]) > 0
    168         assert isinstance(body["created_at"], int)
    169 
    170     async def test_upload_file_exists_on_disk(
    171         self, client: AsyncClient,
    172     ) -> None:
    173         """Uploaded file is actually written to disk."""
    174         content = b"disk test"
    175         resp = await client.post(
    176             "/api/attachments",
    177             files={"file": ("test.txt", content, "text/plain")},
    178             data={"session_id": "sess-1"},
    179         )
    180         assert resp.status_code == 200
    181         body = resp.json()
    182         attachment_id = body["id"]
    183         # Verify file exists by downloading it
    184         dl_resp = await client.get(f"/api/attachments/{attachment_id}")
    185         assert dl_resp.status_code == 200
    186         assert dl_resp.content == content
    187 
    188 
    189 # ---------------------------------------------------------------------------
    190 # Tests — Download
    191 # ---------------------------------------------------------------------------
    192 
    193 class TestDownloadAttachment:
    194     """GET /api/attachments/{id} — file download."""
    195 
    196     async def test_download_returns_file_with_content_type(
    197         self, client: AsyncClient,
    198     ) -> None:
    199         """Download returns the original file with the correct Content-Type."""
    200         png_data = _make_image_png_bytes()
    201         upload_resp = await client.post(
    202             "/api/attachments",
    203             files={"file": ("photo.png", png_data, "image/png")},
    204             data={"session_id": "sess-1"},
    205         )
    206         attachment_id = upload_resp.json()["id"]
    207 
    208         dl_resp = await client.get(f"/api/attachments/{attachment_id}")
    209         assert dl_resp.status_code == 200
    210         assert dl_resp.headers["content-type"] == "image/png"
    211         assert dl_resp.content == png_data
    212 
    213     async def test_download_nonexistent_returns_404(
    214         self, client: AsyncClient,
    215     ) -> None:
    216         """Downloading a non-existent attachment returns 404."""
    217         resp = await client.get("/api/attachments/nonexistent-id")
    218         assert resp.status_code == 404
    219         body = resp.json()
    220         assert "detail" in body
    221         assert "not found" in body["detail"].lower()
    222 
    223 
    224 # ---------------------------------------------------------------------------
    225 # Tests — Image data URL conversion
    226 # ---------------------------------------------------------------------------
    227 
    228 class TestImageDataUrl:
    229     """Images generate data URLs for Hermes consumption."""
    230 
    231     async def test_image_upload_produces_data_url(
    232         self, client: AsyncClient,
    233     ) -> None:
    234         """Image files generate a base64 data URL."""
    235         png_data = _make_image_png_bytes()
    236         resp = await client.post(
    237             "/api/attachments",
    238             files={"file": ("photo.png", png_data, "image/png")},
    239             data={"session_id": "sess-1", "return_data_url": "true"},
    240         )
    241         assert resp.status_code == 200
    242         body = resp.json()
    243         data_url = body.get("data_url", "")
    244         assert data_url.startswith("data:image/png;base64,")
    245         # Verify the base64 part is valid
    246         prefix = "data:image/png;base64,"
    247         encoded_part = data_url[len(prefix):]
    248         decoded = base64.b64decode(encoded_part)
    249         assert decoded == png_data
    250 
    251 
    252 class TestEnsureDb:
    253     """_ensure_db returns db or raises."""
    254 
    255     def test_raises_when_db_is_none(self, monkeypatch) -> None:
    256         import app.routes.attachments as mod
    257         monkeypatch.setattr(mod, "_db", None)
    258         with pytest.raises(RuntimeError, match="not initialized"):
    259             mod._ensure_db()
    260 
    261     def test_returns_db_when_set(self, monkeypatch) -> None:
    262         import app.routes.attachments as mod
    263         from unittest.mock import MagicMock
    264         mock_db = MagicMock()
    265         monkeypatch.setattr(mod, "_db", mock_db)
    266         assert mod._ensure_db() is mock_db
    267 
    268 
    269 class TestAttachmentDir:
    270     """_attachment_dir creates and returns path."""
    271 
    272     def test_creates_directory(self, tmp_path, monkeypatch) -> None:
    273         from app.routes.attachments import _attachment_dir
    274         from unittest.mock import MagicMock
    275         cfg = MagicMock()
    276         cfg.data_dir = str(tmp_path)
    277         monkeypatch.setattr("app.routes.attachments.get_config", lambda: cfg)
    278         result = _attachment_dir("s1")
    279         assert result.exists()
    280         assert result.name == "s1"
    281         assert "attachments" in str(result)
    282         assert "s1" in str(result)
    283 
    284 
    285 class TestSaveFile:
    286     """_save_file writes to disk and returns id + path."""
    287 
    288     def test_writes_file_and_returns_tuple(self, tmp_path, monkeypatch) -> None:
    289         from app.routes.attachments import _save_file
    290         from unittest.mock import MagicMock
    291         cfg = MagicMock()
    292         cfg.data_dir = str(tmp_path)
    293         monkeypatch.setattr("app.routes.attachments.get_config", lambda: cfg)
    294         att_id, path = _save_file(b"hello", "test.txt", "s1")
    295         assert len(att_id) == 36  # UUID
    296         assert path.endswith(".txt")
    297         from pathlib import Path
    298         assert Path(path).read_bytes() == b"hello"
    299 
    300     def test_preserves_extension(self, tmp_path, monkeypatch) -> None:
    301         from app.routes.attachments import _save_file
    302         from unittest.mock import MagicMock
    303         cfg = MagicMock()
    304         cfg.data_dir = str(tmp_path)
    305         monkeypatch.setattr("app.routes.attachments.get_config", lambda: cfg)
    306         _, path = _save_file(b"data", "image.png", "s1")
    307         assert path.endswith(".png")
    308 
    309     def test_empty_extension_when_none(self, tmp_path, monkeypatch) -> None:
    310         from app.routes.attachments import _save_file
    311         from unittest.mock import MagicMock
    312         cfg = MagicMock()
    313         cfg.data_dir = str(tmp_path)
    314         monkeypatch.setattr("app.routes.attachments.get_config", lambda: cfg)
    315         _, path = _save_file(b"data", "noext", "s1")
    316         assert not path.endswith(".")
    317 
    318 
    319 class TestIsImage:
    320     """_is_image checks MIME prefix."""
    321 
    322     def test_image_png_is_image(self) -> None:
    323         from app.routes.attachments import _is_image
    324         assert _is_image("image/png") is True
    325 
    326     def test_image_jpeg_is_image(self) -> None:
    327         from app.routes.attachments import _is_image
    328         assert _is_image("image/jpeg") is True
    329 
    330     def test_pdf_is_not_image(self) -> None:
    331         from app.routes.attachments import _is_image
    332         assert _is_image("application/pdf") is False
    333 
    334     def test_text_is_not_image(self) -> None:
    335         from app.routes.attachments import _is_image
    336         assert _is_image("text/plain") is False
    337 
    338 
    339 class TestToDataUrl:
    340     """_to_data_url encodes bytes as base64 data URL."""
    341 
    342     def test_encodes_png(self) -> None:
    343         from app.routes.attachments import _to_data_url
    344         result = _to_data_url(b"\x89PNG", "image/png")
    345         assert result.startswith("data:image/png;base64,")
    346         # Verify the base64 part is valid
    347         prefix = "data:image/png;base64,"
    348         encoded_part = result[len(prefix):]
    349         decoded = base64.b64decode(encoded_part)
    350         assert decoded == b"\x89PNG"
    351 
    352     def test_encodes_empty(self) -> None:
    353         from app.routes.attachments import _to_data_url
    354         result = _to_data_url(b"", "image/png")
    355         assert result == "data:image/png;base64,"