omnigent-ai--omnigent/tests/server/routes/test_sessions_upload_limits.py

"""Attachment upload type/size enforcement on POST /v1/sessions/{id}/resources/files."""

from __future__ import annotations

from collections.abc import Iterator

import pytest
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse
from fastapi.testclient import TestClient

from omnigent.errors import OmnigentError
from omnigent.runtime.content_resolver import (
    MAX_IMAGE_UPLOAD_BYTES,
    MAX_TEXT_UPLOAD_BYTES,
)
from omnigent.server.routes.sessions import create_sessions_router
from omnigent.stores.agent_store.sqlalchemy_store import SqlAlchemyAgentStore
from omnigent.stores.artifact_store.local import LocalArtifactStore
from omnigent.stores.conversation_store.sqlalchemy_store import (
    SqlAlchemyConversationStore,
)
from omnigent.stores.file_store.sqlalchemy_store import SqlAlchemyFileStore


@pytest.fixture
def upload_client(db_uri: str, tmp_path) -> Iterator[tuple[TestClient, str]]:
    """A sessions route client with file + artifact stores and one session."""
    conversation_store = SqlAlchemyConversationStore(db_uri)
    agent_store = SqlAlchemyAgentStore(db_uri)
    file_store = SqlAlchemyFileStore(db_uri)
    artifact_store = LocalArtifactStore(str(tmp_path / "artifacts"))
    agent_store.create(
        agent_id="ag_test",
        name="test-agent",
        bundle_location="ag_test/bundle",
    )
    conv = conversation_store.create_conversation(title="upload session", agent_id="ag_test")

    app = FastAPI()

    @app.exception_handler(OmnigentError)
    async def _handle_omnigent_error(request: Request, exc: OmnigentError) -> JSONResponse:
        del request
        return JSONResponse(
            status_code=exc.http_status,
            content={"error": {"code": exc.code, "message": exc.message}},
        )

    app.include_router(
        create_sessions_router(
            conversation_store=conversation_store,
            agent_store=agent_store,
            file_store=file_store,
            artifact_store=artifact_store,
        ),
        prefix="/v1",
    )

    with TestClient(app) as client:
        yield client, conv.id


def test_upload_small_text_file_succeeds(upload_client: tuple[TestClient, str]) -> None:
    """A small text file uploads and returns a resource."""
    client, session_id = upload_client
    resp = client.post(
        f"/v1/sessions/{session_id}/resources/files",
        files={"file": ("notes.txt", b"hello world", "text/plain")},
    )
    assert resp.status_code in (200, 201), resp.text
    body = resp.json()
    assert body["name"] == "notes.txt"


def test_upload_rejects_unsupported_type(upload_client: tuple[TestClient, str]) -> None:
    """A pptx (binary office doc) is rejected with 415, not stored."""
    client, session_id = upload_client
    pptx_mime = "application/vnd.openxmlformats-officedocument.presentationml.presentation"
    resp = client.post(
        f"/v1/sessions/{session_id}/resources/files",
        files={"file": ("deck.pptx", b"PK\x03\x04 fake pptx bytes", pptx_mime)},
    )
    assert resp.status_code == 415, resp.text
    assert "Unsupported attachment type" in resp.text


def test_upload_rejects_oversized_image(upload_client: tuple[TestClient, str]) -> None:
    """An image over the per-type limit is rejected with 413."""
    client, session_id = upload_client
    oversized = b"\x00" * (MAX_IMAGE_UPLOAD_BYTES + 1)
    resp = client.post(
        f"/v1/sessions/{session_id}/resources/files",
        files={"file": ("huge.png", oversized, "image/png")},
    )
    assert resp.status_code == 413, resp.status_code
    assert "limit" in resp.text.lower()


def test_upload_csv_mislabeled_as_excel_is_accepted(
    upload_client: tuple[TestClient, str],
) -> None:
    """A .csv the browser tags application/vnd.ms-excel is accepted via the
    extension fallback and stored as a text type (parity with the web client)."""
    client, session_id = upload_client
    resp = client.post(
        f"/v1/sessions/{session_id}/resources/files",
        files={"file": ("data.csv", b"a,b,c\n1,2,3\n", "application/vnd.ms-excel")},
    )
    assert resp.status_code in (200, 201), resp.text
    assert resp.json()["name"] == "data.csv"


def test_upload_text_just_under_limit_succeeds(upload_client: tuple[TestClient, str]) -> None:
    """A text file just under the text cap is accepted."""
    client, session_id = upload_client
    payload = b"a" * (MAX_TEXT_UPLOAD_BYTES - 1024)
    resp = client.post(
        f"/v1/sessions/{session_id}/resources/files",
        files={"file": ("big.txt", payload, "text/plain")},
    )
    assert resp.status_code in (200, 201), resp.status_code


class _FakeUpload:
    """Minimal UploadFile stand-in exposing the chunked ``read`` interface."""

    def __init__(self, data: bytes) -> None:
        self._data = data
        self._pos = 0

    async def read(self, size: int) -> bytes:
        chunk = self._data[self._pos : self._pos + size]
        self._pos += len(chunk)
        return chunk


async def test_read_upload_capped_allows_exactly_at_limit() -> None:
    """A payload exactly at the limit is accepted (the ``>`` boundary)."""
    from omnigent.server.routes.sessions import _read_upload_capped

    data = b"x" * 100
    assert await _read_upload_capped(_FakeUpload(data), 100) == data


async def test_read_upload_capped_rejects_one_over_limit() -> None:
    """One byte over the limit raises HTTP 413."""
    import pytest as _pytest
    from fastapi import HTTPException

    from omnigent.server.routes.sessions import _read_upload_capped

    with _pytest.raises(HTTPException) as exc_info:
        await _read_upload_capped(_FakeUpload(b"x" * 101), 100)
    assert exc_info.value.status_code == 413