mirror of
https://github.com/logos-storage/bittorrent-benchmarks.git
synced 2026-01-02 21:13:11 +00:00
79 lines
2.3 KiB
Python
79 lines
2.3 KiB
Python
import pytest
|
|
from fastapi import FastAPI
|
|
from httpx import AsyncClient, ASGITransport
|
|
from starlette.testclient import TestClient
|
|
|
|
from benchmarks.codex.agent import api
|
|
from benchmarks.codex.agent.agent import CodexAgent
|
|
from benchmarks.codex.agent.tests.fake_codex import FakeCodex
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_create_file():
|
|
codex_client = FakeCodex()
|
|
codex_agent = CodexAgent(codex_client)
|
|
|
|
app = FastAPI()
|
|
app.include_router(api.router)
|
|
app.dependency_overrides[api.codex_agent] = lambda: codex_agent
|
|
|
|
client = TestClient(app)
|
|
|
|
response = client.post(
|
|
"/api/v1/codex/dataset",
|
|
params={"name": "dataset-1", "size": 1024, "seed": 12},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.headers["content-type"] == "text/plain"
|
|
assert response.charset_encoding == "utf-8"
|
|
|
|
manifest = await codex_client.manifest(response.text)
|
|
|
|
assert manifest.datasetSize == 1024
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_should_report_when_download_is_complete():
|
|
codex_client = FakeCodex()
|
|
codex_agent = CodexAgent(codex_client)
|
|
|
|
app = FastAPI()
|
|
app.include_router(api.router)
|
|
app.dependency_overrides[api.codex_agent] = lambda: codex_agent
|
|
|
|
async with AsyncClient(
|
|
transport=ASGITransport(app=app), base_url="http://testserver"
|
|
) as client:
|
|
response = await client.post(
|
|
"/api/v1/codex/dataset",
|
|
params={"name": "dataset-1", "size": 1024, "seed": 12},
|
|
)
|
|
|
|
assert response.status_code == 200
|
|
assert response.charset_encoding == "utf-8"
|
|
|
|
cid = response.text
|
|
|
|
download_stream = codex_client.create_download_stream(cid)
|
|
|
|
response = await client.post(
|
|
"/api/v1/codex/download",
|
|
params={"cid": cid},
|
|
)
|
|
|
|
assert response.status_code == 202
|
|
assert response.json() == {
|
|
"status": f"http://testserver/api/v1/codex/download/{cid}/status"
|
|
}
|
|
|
|
download_stream.feed_data(b"0" * 1024)
|
|
download_stream.feed_eof()
|
|
|
|
await codex_agent.ongoing_downloads[cid].download_task
|
|
|
|
response = await client.get(f"api/v1/codex/download/{cid}/status")
|
|
|
|
assert response.status_code == 200
|
|
assert response.json() == {"downloaded": 1024, "total": 1024}
|