From 1886f10e1904cdd5ba6f41b29dc90ae50096c460 Mon Sep 17 00:00:00 2001
From: bigcat88 <bigcat88@icloud.com>
Date: Tue, 9 Sep 2025 19:30:58 +0300
Subject: [PATCH] add download tests

---
 tests-assets/conftest.py       |  9 ++++---
 tests-assets/test_downloads.py | 45 ++++++++++++++++++++++++++++++++++
 2 files changed, 51 insertions(+), 3 deletions(-)
 create mode 100644 tests-assets/test_downloads.py

diff --git a/tests-assets/conftest.py b/tests-assets/conftest.py
index 24eed0728..16a43d56d 100644
--- a/tests-assets/conftest.py
+++ b/tests-assets/conftest.py
@@ -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,13 +191,16 @@ 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"
-    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}
     form = aiohttp.FormData()
     form.add_field("file", b"A" * 4096, filename=name, content_type="application/octet-stream")
diff --git a/tests-assets/test_downloads.py b/tests-assets/test_downloads.py
new file mode 100644
index 000000000..fad8e2cbb
--- /dev/null
+++ b/tests-assets/test_downloads.py
@@ -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"