feat: add CivitAI stats fetch service

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
2026-03-15 14:32:52 +01:00
parent a8308874d2
commit c1e124bcc8
2 changed files with 263 additions and 0 deletions
@@ -0,0 +1,121 @@
+"""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