feat: enrich list API responses with CivitAI stats from local DB

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-15 14:34:51 +01:00
parent c1e124bcc8
commit c3cd35add7
3 changed files with 113 additions and 4 deletions
+41
View File
@@ -0,0 +1,41 @@
"""Tests for proxy response enrichment with stats."""
import pytest
@pytest.fixture
def db(tmp_path):
from stats_db import StatsDB
db = StatsDB(tmp_path / "test.db")
db.init()
return db
def test_enrich_list_response(db):
"""Stats are merged into list response items."""
from proxy import _enrich_list_response
db.upsert("hash_a", {"download_count": 5000, "rating": 4.8,
"rating_count": 120, "thumbs_up_count": 89})
response_data = {
"items": [
{"sha256": "hash_a", "file_name": "lora_a.safetensors"},
{"sha256": "hash_b", "file_name": "lora_b.safetensors"},
],
"total": 2,
}
enriched = _enrich_list_response(response_data, db)
assert enriched["items"][0]["download_count"] == 5000
assert enriched["items"][0]["rating"] == 4.8
assert enriched["items"][0]["thumbs_up_count"] == 89
# Item without stats gets no extra fields
assert "download_count" not in enriched["items"][1]
def test_enrich_empty_items(db):
"""Empty items list returns unchanged."""
from proxy import _enrich_list_response
response_data = {"items": [], "total": 0}
enriched = _enrich_list_response(response_data, db)
assert enriched["items"] == []