42 lines
1.2 KiB
Python
42 lines
1.2 KiB
Python
"""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"] == []
|