mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-06 14:10:13 -03:00
feat(api): add cursor-based pagination to civitai user-models endpoint
This commit is contained in:
@@ -900,18 +900,28 @@ class FakeMetadataProvider:
|
||||
async def get_model_versions(self, _model_id):
|
||||
return {"modelVersions": [], "name": "", "type": "lora"}
|
||||
|
||||
async def get_user_models(self, _username):
|
||||
return []
|
||||
async def get_user_models(self, _username, cursor=None):
|
||||
return {"items": [], "nextCursor": None}
|
||||
|
||||
async def get_creator_model_count(self, _username):
|
||||
return None
|
||||
|
||||
|
||||
class FakeUserModelsProvider(FakeMetadataProvider):
|
||||
def __init__(self, models):
|
||||
def __init__(self, models, next_cursor=None, estimated_total=None):
|
||||
self.models = models
|
||||
self.next_cursor = next_cursor
|
||||
self.estimated_total = estimated_total
|
||||
self.received_usernames: list[str] = []
|
||||
self.received_cursors: list = []
|
||||
|
||||
async def get_user_models(self, username):
|
||||
async def get_user_models(self, username, cursor=None):
|
||||
self.received_usernames.append(username)
|
||||
return self.models
|
||||
self.received_cursors.append(cursor)
|
||||
return {"items": self.models, "nextCursor": self.next_cursor}
|
||||
|
||||
async def get_creator_model_count(self, _username):
|
||||
return self.estimated_total
|
||||
|
||||
|
||||
async def fake_metadata_provider_factory():
|
||||
@@ -1286,6 +1296,88 @@ async def test_get_civitai_user_models_requires_username():
|
||||
assert "username" in payload["error"].lower()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_civitai_user_models_returns_pagination_fields():
|
||||
models = [
|
||||
{
|
||||
"id": 1,
|
||||
"name": "Model A",
|
||||
"type": "LORA",
|
||||
"tags": [],
|
||||
"modelVersions": [
|
||||
{"id": 100, "name": "v1", "images": [{"url": "http://example.com/a.jpg"}]},
|
||||
],
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"name": "Unsupported",
|
||||
"type": "Other",
|
||||
"modelVersions": [{"id": 200, "name": "v1"}],
|
||||
},
|
||||
]
|
||||
|
||||
provider = FakeUserModelsProvider(models, next_cursor="cursor-token", estimated_total=2140)
|
||||
|
||||
async def provider_factory():
|
||||
return provider
|
||||
|
||||
handler = ModelLibraryHandler(
|
||||
ServiceRegistryAdapter(
|
||||
get_lora_scanner=fake_scanner_factory,
|
||||
get_checkpoint_scanner=fake_scanner_factory,
|
||||
get_embedding_scanner=fake_scanner_factory,
|
||||
get_downloaded_version_history_service=fake_download_history_service_factory,
|
||||
),
|
||||
metadata_provider_factory=provider_factory,
|
||||
)
|
||||
|
||||
response = await handler.get_civitai_user_models(
|
||||
FakeRequest(query={"username": "pixel"})
|
||||
)
|
||||
payload = json.loads(response.text)
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
# modelCount only counts models surviving the type filter
|
||||
assert payload["modelCount"] == 1
|
||||
assert payload["nextCursor"] == "cursor-token"
|
||||
assert payload["hasMore"] is True
|
||||
# first page includes the estimated total
|
||||
assert payload["estimatedTotal"] == 2140
|
||||
assert provider.received_cursors == [None]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_civitai_user_models_passes_cursor_and_omits_estimate():
|
||||
provider = FakeUserModelsProvider([], next_cursor=None, estimated_total=999)
|
||||
|
||||
async def provider_factory():
|
||||
return provider
|
||||
|
||||
handler = ModelLibraryHandler(
|
||||
ServiceRegistryAdapter(
|
||||
get_lora_scanner=fake_scanner_factory,
|
||||
get_checkpoint_scanner=fake_scanner_factory,
|
||||
get_embedding_scanner=fake_scanner_factory,
|
||||
get_downloaded_version_history_service=fake_download_history_service_factory,
|
||||
),
|
||||
metadata_provider_factory=provider_factory,
|
||||
)
|
||||
|
||||
response = await handler.get_civitai_user_models(
|
||||
FakeRequest(query={"username": "pixel", "cursor": "opaque-token"})
|
||||
)
|
||||
payload = json.loads(response.text)
|
||||
|
||||
assert response.status == 200
|
||||
assert payload["success"] is True
|
||||
assert payload["nextCursor"] is None
|
||||
assert payload["hasMore"] is False
|
||||
# cursor requests must not include the estimated total
|
||||
assert payload["estimatedTotal"] is None
|
||||
assert provider.received_cursors == ["opaque-token"]
|
||||
|
||||
|
||||
def test_ensure_handler_mapping_caches_result():
|
||||
call_records = []
|
||||
|
||||
|
||||
@@ -35,9 +35,11 @@ class DummyDownloader:
|
||||
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
|
||||
@@ -622,3 +624,162 @@ async def test_get_image_info_handles_invalid_id(monkeypatch, downloader, caplog
|
||||
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user