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:
@@ -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"] == []
|
||||
Reference in New Issue
Block a user