mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-10-05 01:25:33 -03:00
Will Miao
ec5fef512b
End-to-end verification against the live site found the price capture broken for a whole class of users: the civitai page hosts are not interchangeable, and the user's civitai_host preference was silently fatal. With civitai_host=civitai.red the update DB held zero prices even with tracking enabled. - civitai.red refuses non-browser HTTP clients outright (Cloudflare challenge, 403 for any User-Agent, aiohttp and httpx alike), while civitai.com and civitai.green answer normally for anonymously visible models and 404 for mature ones. An earlier manual check with curl passed on TLS fingerprint luck, which is why this was missed. - get_model_prices now tries the configured host first, then the others, and takes the first parseable payload. The host that worked is remembered, and a host that refuses outright is parked for 15 minutes so a library full of mature models does not pay three requests each; a 404 is model-specific and does not park the host. Links keep using the configured host, which is where the user's own browser has clearance. - Mature models still have no price source anywhere, so that is now stated instead of silent: price_check_attempted_at separates "tried and unreadable" from "never looked", gated versions show a muted "Price unavailable" badge, and the alerts panel reports unavailableCount. - Failures are logged at warning level, once per host per TTL, with the per-host reason, instead of only at debug level. - The recorded alternatives (internal tRPC with the user's API key, or an extension-assisted fetch from the user's browser) and the strengthened upstream ask for a public price field are documented in the plan.
1026 lines
33 KiB
Python
1026 lines
33 KiB
Python
import copy
|
|
import pathlib
|
|
from unittest.mock import AsyncMock
|
|
|
|
import pytest
|
|
|
|
from py.services import civitai_client as civitai_client_module
|
|
from py.services.civitai_client import CivitaiClient
|
|
from py.services.connectivity_guard import OFFLINE_COOLDOWN_ERROR, OFFLINE_FRIENDLY_MESSAGE
|
|
from py.services.errors import RateLimitError, ResourceNotFoundError
|
|
from py.services.model_metadata_provider import ModelMetadataProviderManager
|
|
|
|
|
|
class DummyDownloader:
|
|
def __init__(self):
|
|
self.download_calls = []
|
|
self.memory_calls = []
|
|
self.request_calls = []
|
|
|
|
async def download_file(self, **kwargs):
|
|
self.download_calls.append(kwargs)
|
|
return True, kwargs["save_path"]
|
|
|
|
async def download_to_memory(self, url, use_auth=False):
|
|
self.memory_calls.append({"url": url, "use_auth": use_auth})
|
|
return True, b"bytes", {"content-type": "image/jpeg"}
|
|
|
|
async def make_request(self, method, url, use_auth=True, **kwargs):
|
|
self.request_calls.append(
|
|
{"method": method, "url": url, "use_auth": use_auth, "kwargs": kwargs}
|
|
)
|
|
return True, {}
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def reset_singletons():
|
|
CivitaiClient._instance = None
|
|
ModelMetadataProviderManager._instance = None
|
|
civitai_client_module._creator_model_count_cache.clear()
|
|
yield
|
|
CivitaiClient._instance = None
|
|
ModelMetadataProviderManager._instance = None
|
|
civitai_client_module._creator_model_count_cache.clear()
|
|
|
|
|
|
@pytest.fixture
|
|
def downloader(monkeypatch):
|
|
instance = DummyDownloader()
|
|
monkeypatch.setattr(civitai_client_module, "get_downloader", AsyncMock(return_value=instance))
|
|
return instance
|
|
|
|
|
|
async def test_download_file_uses_downloader(tmp_path, downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
save_dir = tmp_path / "files"
|
|
save_dir.mkdir()
|
|
|
|
success, path = await client.download_file(
|
|
url="https://example.invalid/model",
|
|
save_dir=str(save_dir),
|
|
default_filename="model.safetensors",
|
|
)
|
|
|
|
assert success is True
|
|
assert path == str(save_dir / "model.safetensors")
|
|
assert downloader.download_calls[0]["use_auth"] is True
|
|
|
|
|
|
async def test_client_defaults_to_red_api_host(downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
assert client.base_url == "https://civitai.red/api/v1"
|
|
|
|
|
|
async def test_get_model_by_hash_enriches_metadata(monkeypatch, downloader):
|
|
version_payload = {
|
|
"modelId": 123,
|
|
"model": {"description": "", "tags": []},
|
|
"creator": {},
|
|
"images": [
|
|
{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}},
|
|
{"meta": "not-a-dict"},
|
|
],
|
|
}
|
|
model_payload = {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}}
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
if url.endswith("by-hash/hash"):
|
|
return True, version_payload.copy()
|
|
if url.endswith("/models/123"):
|
|
return True, model_payload
|
|
return False, "unexpected"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_by_hash("hash")
|
|
|
|
assert error is None
|
|
assert result is not None
|
|
assert result["model"]["description"] == "desc"
|
|
assert result["model"]["tags"] == ["tag"]
|
|
assert result["creator"] == {"username": "user"}
|
|
assert "comfy" not in result["images"][0]["meta"]
|
|
assert result["images"][0]["meta"]["other"] == "keep"
|
|
|
|
|
|
async def test_get_model_by_hash_handles_not_found(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, "not found"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_by_hash("missing")
|
|
|
|
assert result is None
|
|
assert error == "Model not found"
|
|
|
|
|
|
async def test_get_model_by_hash_handles_offline_cooldown(downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, OFFLINE_COOLDOWN_ERROR
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_by_hash("missing")
|
|
|
|
assert result is None
|
|
assert error == OFFLINE_FRIENDLY_MESSAGE
|
|
|
|
|
|
async def test_get_model_by_hash_propagates_rate_limit(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, RateLimitError("limited", retry_after=4)
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
with pytest.raises(RateLimitError) as exc_info:
|
|
await client.get_model_by_hash("limited")
|
|
|
|
assert exc_info.value.retry_after == 4
|
|
assert exc_info.value.provider == "civitai_api"
|
|
|
|
|
|
async def test_download_preview_image_writes_file(tmp_path, downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
target = tmp_path / "preview" / "image.jpg"
|
|
|
|
success = await client.download_preview_image("https://example.invalid/preview", str(target))
|
|
|
|
assert success is True
|
|
assert target.exists()
|
|
assert target.read_bytes() == b"bytes"
|
|
|
|
|
|
async def test_download_preview_image_failure(monkeypatch, downloader):
|
|
async def failing_download(url, use_auth=False):
|
|
return False, b"", {}
|
|
|
|
downloader.download_to_memory = failing_download
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
target = "/tmp/ignored.jpg"
|
|
|
|
success = await client.download_preview_image("https://example.invalid/preview", target)
|
|
|
|
assert success is False
|
|
|
|
|
|
async def test_get_model_versions_success(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, {"modelVersions": [{"id": 1}], "type": "LORA", "name": "Model"}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_versions("123")
|
|
|
|
assert result == {"modelVersions": [{"id": 1}], "type": "LORA", "name": "Model"}
|
|
|
|
|
|
async def test_get_model_versions_raises_on_not_found(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, {"message": "Resource not found"}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
with pytest.raises(ResourceNotFoundError):
|
|
await client.get_model_versions("missing")
|
|
|
|
|
|
async def test_get_model_versions_raises_on_nested_not_found(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, {"error": {"message": "Resource not found"}}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
with pytest.raises(ResourceNotFoundError):
|
|
await client.get_model_versions("missing")
|
|
|
|
|
|
async def test_get_model_versions_raises_on_other_errors(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, {"error": {"message": "Server error"}}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
with pytest.raises(RuntimeError):
|
|
await client.get_model_versions("oops")
|
|
|
|
|
|
async def test_get_model_versions_bulk_success(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
assert url.endswith("/models")
|
|
assert kwargs.get("params") == {"ids": "1,2", "nsfw": "true"}
|
|
return True, {
|
|
"items": [
|
|
{
|
|
"id": 1,
|
|
"modelVersions": [{"id": 11}],
|
|
"type": "LORA",
|
|
"name": "One",
|
|
"allowNoCredit": True,
|
|
"allowCommercialUse": ["Sell"],
|
|
"allowDerivatives": True,
|
|
"allowDifferentLicense": True,
|
|
},
|
|
{
|
|
"id": 2,
|
|
"modelVersions": [],
|
|
"type": "Checkpoint",
|
|
"name": "Two",
|
|
"allowNoCredit": False,
|
|
"allowCommercialUse": ["Image"],
|
|
"allowDerivatives": False,
|
|
"allowDifferentLicense": False,
|
|
},
|
|
]
|
|
}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_versions_bulk([1, "2", 2]) # pyright: ignore[reportArgumentType]
|
|
|
|
assert result == {
|
|
1: {
|
|
"modelVersions": [{"id": 11}],
|
|
"type": "LORA",
|
|
"name": "One",
|
|
"allowNoCredit": True,
|
|
"allowCommercialUse": ["Sell"],
|
|
"allowDerivatives": True,
|
|
"allowDifferentLicense": True,
|
|
},
|
|
2: {
|
|
"modelVersions": [],
|
|
"type": "Checkpoint",
|
|
"name": "Two",
|
|
"allowNoCredit": False,
|
|
"allowCommercialUse": ["Image"],
|
|
"allowDerivatives": False,
|
|
"allowDifferentLicense": False,
|
|
},
|
|
}
|
|
|
|
|
|
async def test_get_model_versions_bulk_handles_errors(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, "error"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_versions_bulk([1, 2])
|
|
|
|
assert result is None
|
|
|
|
|
|
async def test_get_model_version_by_version_id(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
if url.endswith("/model-versions/7"):
|
|
return True, {
|
|
"modelId": 321,
|
|
"model": {"description": ""},
|
|
"files": [],
|
|
"images": [{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}}],
|
|
}
|
|
if url.endswith("/models/321"):
|
|
return True, {"description": "desc", "tags": ["tag"], "creator": {"username": "user"}}
|
|
return False, "unexpected"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_version(version_id=7)
|
|
|
|
assert result is not None
|
|
assert result["model"]["description"] == "desc"
|
|
assert result["model"]["tags"] == ["tag"]
|
|
assert result["creator"] == {"username": "user"}
|
|
assert "comfy" not in result["images"][0]["meta"]
|
|
assert result["images"][0]["meta"]["other"] == "keep"
|
|
|
|
|
|
async def test_get_model_version_with_model_id_prefers_version_endpoint(monkeypatch, downloader):
|
|
requests = []
|
|
|
|
model_payload = {
|
|
"modelVersions": [
|
|
{
|
|
"id": 7,
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"hashes": {"SHA256": "hash7"},
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"description": "desc",
|
|
"tags": ["tag"],
|
|
"creator": {"username": "user"},
|
|
"name": "Model",
|
|
"type": "LORA",
|
|
"nsfw": False,
|
|
"poi": False,
|
|
}
|
|
|
|
version_payload = {
|
|
"id": 7,
|
|
"modelId": 99,
|
|
"model": {},
|
|
"files": [],
|
|
"images": [],
|
|
}
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requests.append(url)
|
|
if url.endswith("/models/99"):
|
|
return True, copy.deepcopy(model_payload)
|
|
if url.endswith("/model-versions/7"):
|
|
return True, copy.deepcopy(version_payload)
|
|
return False, "unexpected"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_version(model_id=99, version_id=7)
|
|
|
|
assert result is not None
|
|
assert result["id"] == 7
|
|
assert result["model"]["description"] == "desc"
|
|
assert result["model"]["tags"] == ["tag"]
|
|
assert result["creator"] == {"username": "user"}
|
|
assert requests[0].endswith("/models/99")
|
|
assert requests[1].endswith("/model-versions/7")
|
|
|
|
|
|
async def test_get_model_version_with_model_id_fallbacks_to_hash(monkeypatch, downloader):
|
|
requests = []
|
|
|
|
model_payload = {
|
|
"modelVersions": [
|
|
{
|
|
"id": 7,
|
|
"files": [
|
|
{
|
|
"type": "Model",
|
|
"primary": True,
|
|
"hashes": {"SHA256": "hash7"},
|
|
}
|
|
],
|
|
}
|
|
],
|
|
"description": "desc",
|
|
"tags": ["tag"],
|
|
"creator": {"username": "user"},
|
|
"name": "Model",
|
|
"type": "LORA",
|
|
"nsfw": False,
|
|
"poi": False,
|
|
}
|
|
|
|
version_payload = {
|
|
"id": 7,
|
|
"modelId": 99,
|
|
"files": [],
|
|
"images": [],
|
|
}
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requests.append(url)
|
|
if url.endswith("/models/99"):
|
|
return True, copy.deepcopy(model_payload)
|
|
if url.endswith("/model-versions/7"):
|
|
return False, "boom"
|
|
if url.endswith("/model-versions/by-hash/hash7"):
|
|
return True, copy.deepcopy(version_payload)
|
|
return False, "unexpected"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_version(model_id=99, version_id=7)
|
|
|
|
assert result is not None
|
|
assert result["id"] == 7
|
|
assert result["model"]["description"] == "desc"
|
|
assert result["model"]["tags"] == ["tag"]
|
|
assert result["creator"] == {"username": "user"}
|
|
assert requests[1].endswith("/model-versions/7")
|
|
assert requests[2].endswith("/model-versions/by-hash/hash7")
|
|
|
|
|
|
async def test_get_model_version_with_model_id_builds_from_model_data(monkeypatch, downloader):
|
|
model_payload = {
|
|
"modelVersions": [
|
|
{
|
|
"id": 7,
|
|
"files": [],
|
|
"name": "v1",
|
|
}
|
|
],
|
|
"description": "desc",
|
|
"tags": ["tag"],
|
|
"creator": {"username": "user"},
|
|
"name": "Model",
|
|
"type": "LORA",
|
|
"nsfw": False,
|
|
"poi": False,
|
|
}
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
if url.endswith("/models/99"):
|
|
return True, copy.deepcopy(model_payload)
|
|
if url.endswith("/model-versions/7"):
|
|
return False, "boom"
|
|
if "/model-versions/by-hash/" in url:
|
|
return False, "boom"
|
|
return False, "unexpected"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_model_version(model_id=99, version_id=7)
|
|
|
|
assert result is not None
|
|
assert result["modelId"] == 99
|
|
assert result["model"]["name"] == "Model"
|
|
assert result["model"]["type"] == "LORA"
|
|
assert result["model"]["description"] == "desc"
|
|
assert result["model"]["tags"] == ["tag"]
|
|
assert result["creator"] == {"username": "user"}
|
|
|
|
|
|
async def test_get_model_version_requires_identifier(monkeypatch, downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
result = await client.get_model_version()
|
|
assert result is None
|
|
|
|
|
|
async def test_get_model_version_info_handles_not_found(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, "not found"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_version_info("55")
|
|
|
|
assert result is None
|
|
assert error == "Model not found"
|
|
|
|
|
|
async def test_get_model_version_info_success(monkeypatch, downloader):
|
|
expected = {"id": 55, "images": [{"meta": {"comfy": {"foo": "bar"}, "other": "keep"}}]}
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, expected
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_version_info("55")
|
|
|
|
assert result == expected
|
|
assert error is None
|
|
assert result is not None
|
|
assert "comfy" not in result["images"][0]["meta"]
|
|
assert result["images"][0]["meta"]["other"] == "keep"
|
|
|
|
|
|
async def test_get_image_info_returns_matching_item(monkeypatch, downloader):
|
|
"""When API returns multiple items, return the one matching the requested ID."""
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
# Requested ID is 42, but it's the second item in the response
|
|
return True, {"items": [{"id": 41}, {"id": 42, "name": "target"}, {"id": 43}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info("42")
|
|
|
|
assert result == {"id": 42, "name": "target"}
|
|
|
|
|
|
async def test_get_image_info_returns_none_when_id_mismatch(monkeypatch, downloader, caplog):
|
|
"""When API returns items but none match the requested ID, return None and log warning."""
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
# Requested ID is 999, but API returns different IDs (simulating deleted/hidden image)
|
|
return True, {"items": [{"id": 1}, {"id": 2}, {"id": 3}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info("999")
|
|
|
|
assert result is None
|
|
# Verify warning was logged
|
|
assert "CivitAI API returned no matching image for requested ID 999" in caplog.text
|
|
assert "Returned 3 item(s) with IDs: [1, 2, 3]" in caplog.text
|
|
|
|
|
|
async def test_get_image_info_handles_missing(monkeypatch, downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, {"items": []}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info("42")
|
|
|
|
assert result is None
|
|
|
|
|
|
async def test_get_image_info_prefers_red_host_for_red_source(monkeypatch, downloader):
|
|
requested_urls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requested_urls.append(url)
|
|
return True, {"items": [{"id": 124950237, "name": "target"}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info(
|
|
"124950237", source_url="https://civitai.red/images/124950237"
|
|
)
|
|
|
|
assert result == {"id": 124950237, "name": "target"}
|
|
assert requested_urls == [
|
|
"https://civitai.red/api/v1/images?imageId=124950237&nsfw=X&withMeta=true"
|
|
]
|
|
|
|
|
|
async def test_get_image_info_uses_red_host_even_for_red_source(monkeypatch, downloader):
|
|
requested_urls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requested_urls.append(url)
|
|
return True, {"items": [{"id": 124950237, "name": "target"}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info(
|
|
"124950237", source_url="https://civitai.red/images/124950237"
|
|
)
|
|
|
|
assert result == {"id": 124950237, "name": "target"}
|
|
assert requested_urls == [
|
|
"https://civitai.red/api/v1/images?imageId=124950237&nsfw=X&withMeta=true",
|
|
]
|
|
|
|
|
|
async def test_get_image_info_does_not_fall_back_after_request_failure(monkeypatch, downloader):
|
|
requested_urls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requested_urls.append(url)
|
|
return False, "403 forbidden"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info(
|
|
"124950237", source_url="https://civitai.red/images/124950237"
|
|
)
|
|
|
|
assert result is None
|
|
assert requested_urls == [
|
|
"https://civitai.red/api/v1/images?imageId=124950237&nsfw=X&withMeta=true",
|
|
]
|
|
|
|
|
|
async def test_get_image_info_handles_invalid_id(monkeypatch, downloader, caplog):
|
|
"""When given a non-numeric image ID, return None and log error."""
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result = await client.get_image_info("not-a-number")
|
|
|
|
assert result is None
|
|
assert "Invalid image ID format" in caplog.text
|
|
|
|
|
|
async def test_get_user_models_requests_first_page_with_stable_params(downloader):
|
|
request_calls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
request_calls.append({"method": method, "url": url, "kwargs": kwargs})
|
|
return True, {
|
|
"items": [
|
|
{
|
|
"id": 1,
|
|
"modelVersions": [
|
|
{"id": 100, "images": [{"meta": {"comfy": {"x": 1}}}]}
|
|
],
|
|
}
|
|
],
|
|
"metadata": {"nextCursor": "next-token"},
|
|
}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
result = await client.get_user_models("pixel")
|
|
|
|
assert result is not None
|
|
assert result["nextCursor"] == "next-token"
|
|
assert len(result["items"]) == 1
|
|
# comfy metadata is still stripped
|
|
assert "comfy" not in result["items"][0]["modelVersions"][0]["images"][0]["meta"]
|
|
|
|
call = request_calls[0]
|
|
assert call["method"] == "GET"
|
|
assert call["url"] == "https://civitai.red/api/v1/models"
|
|
params = call["kwargs"]["params"]
|
|
assert params["username"] == "pixel"
|
|
assert params["nsfw"] == "true"
|
|
assert params["limit"] == 100
|
|
assert params["sort"] == "Newest"
|
|
assert params["period"] == "AllTime"
|
|
assert "cursor" not in params
|
|
|
|
|
|
async def test_get_user_models_passes_cursor_and_stringifies_next_cursor(downloader):
|
|
request_calls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
request_calls.append(kwargs)
|
|
return True, {"items": [], "metadata": {"nextCursor": 12345}}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
result = await client.get_user_models("pixel", cursor="opaque-token")
|
|
|
|
assert request_calls[0]["params"]["cursor"] == "opaque-token"
|
|
assert result == {"items": [], "nextCursor": "12345"}
|
|
|
|
|
|
async def test_get_user_models_without_next_cursor_returns_none_cursor(downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, {"items": [{"id": 1, "modelVersions": []}], "metadata": {}}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
result = await client.get_user_models("pixel")
|
|
|
|
assert result == {"items": [{"id": 1, "modelVersions": []}], "nextCursor": None}
|
|
|
|
|
|
async def test_get_user_models_failure_returns_none(downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, "500 server error"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
result = await client.get_user_models("pixel")
|
|
|
|
assert result is None
|
|
|
|
|
|
async def test_get_creator_model_count_matches_exact_username(downloader):
|
|
request_calls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
request_calls.append({"url": url, "kwargs": kwargs})
|
|
return True, {
|
|
"items": [
|
|
{"username": "pixelart", "modelCount": 5},
|
|
{"username": "Pixel", "modelCount": 2140},
|
|
]
|
|
}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
count = await client.get_creator_model_count("pixel")
|
|
|
|
assert count == 2140
|
|
assert request_calls[0]["url"] == "https://civitai.red/api/v1/creators"
|
|
assert request_calls[0]["kwargs"]["params"] == {"query": "pixel", "limit": 10}
|
|
|
|
|
|
async def test_get_creator_model_count_without_exact_match_returns_none(downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, {"items": [{"username": "pixelart", "modelCount": 5}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
count = await client.get_creator_model_count("pixel")
|
|
|
|
assert count is None
|
|
|
|
|
|
async def test_get_creator_model_count_caches_results(downloader):
|
|
request_count = 0
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
nonlocal request_count
|
|
request_count += 1
|
|
return True, {"items": [{"username": "pixel", "modelCount": 42}]}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
assert await client.get_creator_model_count("pixel") == 42
|
|
# case-insensitive cache key, second call served from cache
|
|
assert await client.get_creator_model_count("Pixel") == 42
|
|
assert request_count == 1
|
|
|
|
|
|
async def test_get_creator_model_count_caches_failures(downloader):
|
|
request_count = 0
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
nonlocal request_count
|
|
request_count += 1
|
|
return False, "500 server error"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
assert await client.get_creator_model_count("pixel") is None
|
|
assert await client.get_creator_model_count("pixel") is None
|
|
assert request_count == 1
|
|
|
|
|
|
async def test_get_creator_model_count_never_raises(downloader):
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, "unexpected non-dict payload"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
assert await client.get_creator_model_count("pixel") is None
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"placeholder_hash",
|
|
[
|
|
"e3b0c44298", # AutoV2 (10 chars)
|
|
"e3b0c44298fc", # AutoV3 (12 chars)
|
|
"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", # full SHA256
|
|
],
|
|
)
|
|
async def test_get_model_by_hash_rejects_empty_placeholder_without_request(downloader, placeholder_hash):
|
|
"""The empty-hash placeholder must never be resolved via the by-hash API:
|
|
CivitAI's index can contain polluted entries for it (e.g. a broken SD 1.5
|
|
LoRA whose AutoV3 equals the placeholder)."""
|
|
requested = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
requested.append(url)
|
|
return True, {}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
result, error = await client.get_model_by_hash(placeholder_hash)
|
|
|
|
assert result is None
|
|
assert error == "Model not found"
|
|
assert requested == []
|
|
|
|
|
|
async def test_get_version_file_mini_returns_payload(downloader):
|
|
"""The mini endpoint returns the raw stored filename (#1100)."""
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
assert method == "GET"
|
|
assert url.endswith("/model-versions/mini/3284136")
|
|
assert kwargs.get("params") == {"modelFileId": 3168412}
|
|
assert use_auth is True
|
|
return True, {"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
result = await client.get_version_file_mini(3284136, 3168412)
|
|
|
|
assert result == {"fileName": "CyberRealistic_zit_v8.0_bf16.safetensors"}
|
|
|
|
|
|
async def test_get_version_file_mini_returns_none_on_failure(downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, "Model file 2 not found in version 1"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
assert await client.get_version_file_mini(1, 2) is None
|
|
|
|
|
|
async def test_get_version_file_mini_propagates_rate_limit(downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, RateLimitError("limited", retry_after=1.0)
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
with pytest.raises(RateLimitError):
|
|
await client.get_version_file_mini(1, 2)
|
|
|
|
|
|
async def test_get_model_prices_parses_public_page(downloader):
|
|
"""Prices come from the page payload, read anonymously (no API key)."""
|
|
client = await CivitaiClient.get_instance()
|
|
page_html = (
|
|
'<script id="__NEXT_DATA__" type="application/json">'
|
|
'{"props":{"pageProps":{"trpcState":{"json":{"queries":['
|
|
'{"queryKey":[["model","getById"],{"input":{"id":7}}],"state":{"data":'
|
|
'{"modelVersions":[{"id":42,"paidAccess":{"endsAt":null,'
|
|
'"timeframeDays":null,"terms":{"download":{"price":5000}},'
|
|
'"sale":null}}]}}}]}}}}}</script>'
|
|
)
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
assert method == "GET"
|
|
assert "/models/7" in url
|
|
assert use_auth is False
|
|
assert kwargs.get("custom_headers", {}).get("Accept") == "text/html"
|
|
return True, page_html
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
result = await client.get_model_prices(7)
|
|
|
|
assert result is not None
|
|
assert result[42]["price_buzz"] == 5000
|
|
|
|
|
|
async def test_get_model_prices_returns_none_on_unusable_page(downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, "<html><body>challenge</body></html>"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
assert await client.get_model_prices(7) is None
|
|
|
|
|
|
async def test_get_model_prices_rejects_json_body(downloader):
|
|
"""A JSON response is not the page; it must not be parsed as one."""
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return True, {"error": "nope"}
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
assert await client.get_model_prices(7) is None
|
|
|
|
|
|
async def test_get_model_prices_propagates_rate_limit(downloader):
|
|
client = await CivitaiClient.get_instance()
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
return False, RateLimitError("limited", retry_after=1.0)
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
with pytest.raises(RateLimitError):
|
|
await client.get_model_prices(7)
|
|
|
|
|
|
# --- Model page host fallback -------------------------------------------------
|
|
|
|
_PAGE_FIXTURE = (
|
|
pathlib.Path(__file__).resolve().parents[1]
|
|
/ "utils"
|
|
/ "fixtures"
|
|
/ "civitai_model_page_paid.html"
|
|
)
|
|
|
|
|
|
def _page_html() -> str:
|
|
return _PAGE_FIXTURE.read_text(encoding="utf-8")
|
|
|
|
|
|
async def test_get_model_prices_falls_back_when_the_preferred_host_refuses(
|
|
downloader, monkeypatch
|
|
):
|
|
"""civitai.red refuses non-browser clients (Cloudflare); the price must still
|
|
be readable from a host that answers. This is the user-visible bug: with
|
|
civitai_host=civitai.red every fetch used to fail."""
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
monkeypatch.setattr(client, "_page_host", lambda: "civitai.red")
|
|
seen = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
seen.append(url)
|
|
if "civitai.red" in url:
|
|
return False, "Access forbidden"
|
|
return True, _page_html()
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
result = await client.get_model_prices(4242)
|
|
|
|
assert result is not None
|
|
assert result[1001]["price_buzz"] == 5000
|
|
assert seen[0].startswith("https://civitai.red/")
|
|
assert any("civitai.com" in url for url in seen)
|
|
# The working host is remembered and the refusing one is parked.
|
|
assert client._page_host_preference == "civitai.com"
|
|
assert "civitai.red" in client._page_host_blocked
|
|
|
|
|
|
async def test_get_model_prices_reuses_the_working_host(downloader, monkeypatch):
|
|
client = await CivitaiClient.get_instance()
|
|
monkeypatch.setattr(client, "_page_host", lambda: "civitai.red")
|
|
calls = []
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
calls.append(url)
|
|
if "civitai.red" in url:
|
|
return False, "Access forbidden"
|
|
return True, _page_html()
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
await client.get_model_prices(4242)
|
|
calls.clear()
|
|
await client.get_model_prices(4242)
|
|
|
|
# Second time around the parked host is not retried.
|
|
assert calls and all("civitai.red" not in url for url in calls)
|
|
|
|
|
|
async def test_get_model_prices_returns_none_when_no_host_can_serve(downloader, monkeypatch):
|
|
"""Mature models: 404 on civitai.com/green, challenge on civitai.red."""
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
monkeypatch.setattr(client, "_page_host", lambda: "civitai.com")
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
if "civitai.red" in url:
|
|
return False, "Access forbidden"
|
|
return False, "Resource not found"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
assert await client.get_model_prices(2981320) is None
|
|
|
|
|
|
async def test_get_model_prices_does_not_park_a_host_on_404(downloader, monkeypatch):
|
|
"""A 404 is model-specific (mature content hidden anonymously), not a reason
|
|
to stop using the host for other models."""
|
|
|
|
client = await CivitaiClient.get_instance()
|
|
monkeypatch.setattr(client, "_page_host", lambda: "civitai.com")
|
|
|
|
async def fake_make_request(method, url, use_auth=True, **kwargs):
|
|
if "civitai.com" in url:
|
|
return False, "Resource not found"
|
|
return False, "Access forbidden"
|
|
|
|
downloader.make_request = fake_make_request
|
|
|
|
await client.get_model_prices(1)
|
|
|
|
assert "civitai.com" not in client._page_host_blocked
|
|
assert "civitai.red" in client._page_host_blocked
|