Spaces:
Runtime error
Runtime error
File size: 6,105 Bytes
290ff9e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 | """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
|