122 lines
3.5 KiB
Python
122 lines
3.5 KiB
Python
"""Tests for the CivitAI stats fetch service."""
|
|
import pytest
|
|
from unittest.mock import patch
|
|
|
|
|
|
@pytest.fixture
|
|
def db(tmp_path):
|
|
from stats_db import StatsDB
|
|
db = StatsDB(tmp_path / "test.db")
|
|
db.init()
|
|
return db
|
|
|
|
|
|
@pytest.fixture
|
|
def service(db):
|
|
from stats_service import StatsFetchService
|
|
return StatsFetchService(db)
|
|
|
|
|
|
def test_extract_version_stats():
|
|
"""Extract stats from a CivitAI model API response."""
|
|
from stats_service import extract_version_stats
|
|
|
|
model_response = {
|
|
"id": 123,
|
|
"modelVersions": [
|
|
{
|
|
"id": 456,
|
|
"files": [{"hashes": {"SHA256": "AABB"}}],
|
|
"stats": {
|
|
"downloadCount": 5000,
|
|
"ratingCount": 120,
|
|
"rating": 4.8,
|
|
"thumbsUpCount": 89,
|
|
},
|
|
},
|
|
{
|
|
"id": 789,
|
|
"files": [{"hashes": {"SHA256": "CCDD"}}],
|
|
"stats": {
|
|
"downloadCount": 1000,
|
|
"ratingCount": 30,
|
|
"rating": 3.5,
|
|
"thumbsUpCount": 20,
|
|
},
|
|
},
|
|
],
|
|
}
|
|
result = extract_version_stats(model_response)
|
|
assert len(result) == 2
|
|
assert result[0][0] == "aabb" # sha256 lowercased
|
|
assert result[0][1]["download_count"] == 5000
|
|
assert result[0][1]["civitai_model_id"] == 123
|
|
assert result[0][1]["civitai_version_id"] == 456
|
|
assert result[1][0] == "ccdd"
|
|
|
|
|
|
def test_extract_version_stats_no_hash():
|
|
"""Versions without a SHA256 hash are skipped."""
|
|
from stats_service import extract_version_stats
|
|
|
|
model_response = {
|
|
"id": 1,
|
|
"modelVersions": [
|
|
{"id": 10, "files": [{"hashes": {}}], "stats": {"downloadCount": 1}},
|
|
],
|
|
}
|
|
assert extract_version_stats(model_response) == []
|
|
|
|
|
|
def test_extract_version_stats_no_stats():
|
|
"""Versions without stats get zeros."""
|
|
from stats_service import extract_version_stats
|
|
|
|
model_response = {
|
|
"id": 1,
|
|
"modelVersions": [
|
|
{"id": 10, "files": [{"hashes": {"SHA256": "AABB"}}]},
|
|
],
|
|
}
|
|
result = extract_version_stats(model_response)
|
|
assert len(result) == 1
|
|
assert result[0][1]["download_count"] == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_fetch_stats_for_models(service):
|
|
"""fetch_stats_for_models calls CivitAI API and stores results."""
|
|
models = [
|
|
{"sha256": "aabb", "civitai_model_id": 123},
|
|
{"sha256": "ccdd", "civitai_model_id": 456},
|
|
]
|
|
mock_response_123 = {
|
|
"id": 123,
|
|
"modelVersions": [{
|
|
"id": 10,
|
|
"files": [{"hashes": {"SHA256": "AABB"}}],
|
|
"stats": {"downloadCount": 500, "ratingCount": 10,
|
|
"rating": 4.0, "thumbsUpCount": 8},
|
|
}],
|
|
}
|
|
mock_response_456 = {
|
|
"id": 456,
|
|
"modelVersions": [{
|
|
"id": 20,
|
|
"files": [{"hashes": {"SHA256": "CCDD"}}],
|
|
"stats": {"downloadCount": 100, "ratingCount": 5,
|
|
"rating": 3.0, "thumbsUpCount": 2},
|
|
}],
|
|
}
|
|
|
|
async def mock_fetch(model_id):
|
|
return {123: mock_response_123, 456: mock_response_456}.get(model_id)
|
|
|
|
with patch.object(service, "_fetch_model", side_effect=mock_fetch):
|
|
count = await service.fetch_stats_for_models(models)
|
|
|
|
assert count == 2
|
|
result = service.db.get_by_hashes(["aabb", "ccdd"])
|
|
assert result["aabb"]["download_count"] == 500
|
|
assert result["ccdd"]["download_count"] == 100
|