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 tempfile
import time import time
from pathlib import Path from pathlib import Path
from typing import AsyncIterator, Callable from typing import AsyncIterator, Callable, Optional
import aiohttp import aiohttp
import pytest import pytest
@ -191,13 +191,16 @@ async def asset_factory(http: aiohttp.ClientSession, api_base: str):
@pytest_asyncio.fixture @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>. Upload one asset into models/checkpoints/unit-tests/<name>.
Returns response dict with id, asset_hash, tags, etc. Returns response dict with id, asset_hash, tags, etc.
""" """
name = "unit_1_example.safetensors" name = "unit_1_example.safetensors"
tags = ["models", "checkpoints", "unit-tests", "alpha"] 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} meta = {"purpose": "test", "epoch": 1, "flags": ["x", "y"], "nullable": None}
form = aiohttp.FormData() form = aiohttp.FormData()
form.add_field("file", b"A" * 4096, filename=name, content_type="application/octet-stream") form.add_field("file", b"A" * 4096, filename=name, content_type="application/octet-stream")

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"