add download tests

This commit is contained in:
bigcat88 2025-09-09 19:30:58 +03:00
parent 357193f7b5
commit 1886f10e19
No known key found for this signature in database
GPG Key ID: 1F0BF0EC3CF22721
2 changed files with 51 additions and 3 deletions

View File

@ -7,7 +7,7 @@ import sys
import tempfile
import time
from pathlib import Path
from typing import AsyncIterator, Callable
from typing import AsyncIterator, Callable, Optional
import aiohttp
import pytest
@ -191,12 +191,15 @@ async def asset_factory(http: aiohttp.ClientSession, api_base: str):
@pytest_asyncio.fixture
async def seeded_asset(http: aiohttp.ClientSession, api_base: str) -> dict:
async def seeded_asset(request: pytest.FixtureRequest, http: aiohttp.ClientSession, api_base: str) -> dict:
"""
Upload one asset into models/checkpoints/unit-tests/<name>.
Returns response dict with id, asset_hash, tags, etc.
"""
name = "unit_1_example.safetensors"
p = getattr(request, "param", {}) or {}
tags: Optional[list[str]] = p.get("tags")
if tags is None:
tags = ["models", "checkpoints", "unit-tests", "alpha"]
meta = {"purpose": "test", "epoch": 1, "flags": ["x", "y"], "nullable": None}
form = aiohttp.FormData()

View File

@ -0,0 +1,45 @@
from pathlib import Path
import aiohttp
import pytest
@pytest.mark.asyncio
async def test_download_attachment_and_inline(http: aiohttp.ClientSession, api_base: str, seeded_asset: dict):
aid = seeded_asset["id"]
# default attachment
async with http.get(f"{api_base}/api/assets/{aid}/content") as r1:
data = await r1.read()
assert r1.status == 200
cd = r1.headers.get("Content-Disposition", "")
assert "attachment" in cd
assert data and len(data) == 4096
# inline requested
async with http.get(f"{api_base}/api/assets/{aid}/content?disposition=inline") as r2:
await r2.read()
assert r2.status == 200
cd2 = r2.headers.get("Content-Disposition", "")
assert "inline" in cd2
@pytest.mark.asyncio
@pytest.mark.parametrize("seeded_asset", [{"tags": ["models", "checkpoints"]}], indirect=True)
async def test_download_missing_file_returns_404(
http: aiohttp.ClientSession, api_base: str, comfy_tmp_base_dir: Path, seeded_asset: dict
):
# Remove the underlying file then attempt download.
# We initialize fixture without additional tags to know exactly the asset file path.
aid = seeded_asset["id"]
async with http.get(f"{api_base}/api/assets/{aid}") as rg:
detail = await rg.json()
assert rg.status == 200
rel_inside_category = detail["name"]
abs_path = comfy_tmp_base_dir / "models" / "checkpoints" / rel_inside_category
if abs_path.exists():
abs_path.unlink()
async with http.get(f"{api_base}/api/assets/{aid}/content") as r2:
body = await r2.json()
assert r2.status == 404
assert body["error"]["code"] == "FILE_NOT_FOUND"