tests/test_uploads_api.py

"""Tests for the /uploads/* endpoints and /assess s3_key normalization.

Each test installs a stub asset storage on the live FastAPI app, exercises
the route via TestClient, and asserts the contract the frontend depends on.
"""

from __future__ import annotations

import pytest
from fastapi.testclient import TestClient

from ergo_agentic import server
from ergo_agentic.config import S3Config
from ergo_agentic.storage.s3 import S3AssetStorage


class _FakeS3:
    def __init__(self) -> None:
        self.put_calls: list[dict] = []

    def put_object(self, **kwargs):
        self.put_calls.append(kwargs)

    def upload_fileobj(self, **kwargs):
        self.put_calls.append({"streamed": True, **kwargs})

    def generate_presigned_url(self, op, *, Params, ExpiresIn):  # noqa: N803  (boto3 contract)
        return f"https://signed.example/{op}/{Params['Key']}?ttl={ExpiresIn}"


def _make_storage() -> S3AssetStorage:
    return S3AssetStorage(
        S3Config(
            enabled=True,
            bucket_name="ergo-test",
            asset_prefix="ergo/assessments",
            presigned_url_ttl_seconds=3600,
        ),
        client=_FakeS3(),
    )


@pytest.fixture
def storage_client():
    """A TestClient with the lifespan already started, plus a stub storage
    installed *after* lifespan startup. The lifespan resets _asset_storage
    based on S3_ENABLED env, so we have to install the stub once the app is
    live."""
    with TestClient(server.app) as client:
        previous = server._asset_storage
        server._asset_storage = _make_storage()
        try:
            yield client
        finally:
            server._asset_storage = previous


def test_presign_endpoint_returns_scoped_keys(storage_client):
    res = storage_client.post(
        "/uploads/presign",
        json={
            "run_id": "run_abc",
            "files": [
                {"image_id": "img_1", "filename": "chair.JPG"},
                {
                    "image_id": "img_2",
                    "filename": "desk",
                    "content_type": "image/png",
                },
            ],
        },
    )
    assert res.status_code == 200, res.text
    body = res.json()
    assert body["run_id"] == "run_abc"
    keys = [u["s3_key"] for u in body["uploads"]]
    assert keys[0].endswith("/originals/img_1.jpg")
    assert keys[1].endswith("/originals/img_2.png")
    for u in body["uploads"]:
        assert u["put_url"].startswith("https://signed.example/put_object/")
        assert u["get_url"].startswith("https://signed.example/get_object/")


def test_presign_rejects_duplicate_image_ids(storage_client):
    res = storage_client.post(
        "/uploads/presign",
        json={
            "run_id": "run_abc",
            "files": [
                {"image_id": "img_1"},
                {"image_id": "img_1"},
            ],
        },
    )
    assert res.status_code == 400
    assert "duplicate" in res.json()["detail"].lower()


def test_presign_returns_503_when_storage_disabled():
    with TestClient(server.app) as client:
        # Lifespan already cleared _asset_storage because S3_ENABLED is unset.
        server._asset_storage = None
        res = client.post(
            "/uploads/presign",
            json={"files": [{"image_id": "img_1"}]},
        )
    assert res.status_code == 503


def test_refresh_rejects_keys_outside_run(storage_client):
    res = storage_client.post(
        "/uploads/refresh",
        json={
            "run_id": "run_abc",
            "keys": ["ergo/assessments/run_other/x.jpg"],
        },
    )
    assert res.status_code == 400
    assert "outside run" in res.json()["detail"]


def test_refresh_signs_in_run_keys(storage_client):
    res = storage_client.post(
        "/uploads/refresh",
        json={
            "run_id": "run_abc",
            "keys": ["ergo/assessments/run_abc/originals/img_1.jpg"],
        },
    )
    assert res.status_code == 200
    [u] = res.json()["urls"]
    assert u["s3_key"] == "ergo/assessments/run_abc/originals/img_1.jpg"
    assert u["get_url"].startswith("https://signed.example/get_object/")


def test_to_image_inputs_resolves_s3_key():
    previous = server._asset_storage
    server._asset_storage = _make_storage()
    try:
        inputs = server._to_image_inputs(
            [
                server.ImageRequest(
                    s3_key="ergo/assessments/run_abc/originals/img_1.jpg",
                    image_id="img_1",
                    content_type="image/jpeg",
                )
            ]
        )
    finally:
        server._asset_storage = previous
    assert inputs[0]["s3_key"] == "ergo/assessments/run_abc/originals/img_1.jpg"
    assert inputs[0]["url"].startswith("https://signed.example/get_object/")


def test_to_image_inputs_rejects_s3_key_outside_prefix():
    from fastapi import HTTPException

    previous = server._asset_storage
    server._asset_storage = _make_storage()
    try:
        with pytest.raises(HTTPException) as excinfo:
            server._to_image_inputs(
                [server.ImageRequest(s3_key="other-prefix/x.jpg", image_id="img_1")]
            )
    finally:
        server._asset_storage = previous
    assert excinfo.value.status_code == 400


def test_to_image_inputs_rejects_missing_source():
    from fastapi import HTTPException

    with pytest.raises(HTTPException) as excinfo:
        server._to_image_inputs([server.ImageRequest(image_id="img_1")])
    assert excinfo.value.status_code == 422


def test_to_image_inputs_passes_through_url():
    inputs = server._to_image_inputs(
        [server.ImageRequest(url="https://example.com/x.jpg", image_id="img_1")]
    )
    assert inputs[0]["url"] == "https://example.com/x.jpg"
    assert inputs[0]["s3_key"] is None


def test_info_endpoint_reports_s3_status(storage_client):
    body = storage_client.get("/").json()
    assert body["s3_enabled"] is True