feat: add SQLite stats database layer for CivitAI stats

Implements StatsDB class with upsert, batch upsert, get_by_hashes,
get_all, count, and close methods backed by a WAL-mode SQLite database.
Also guards __init__.py relative imports behind __package__ check so
pytest can run without ComfyUI context, and removes empty tests/__init__.py
to prevent pytest Package collector from traversing into the plugin root.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-03-15 14:29:59 +01:00
parent aefb129e07
commit 44088649d8
5 changed files with 271 additions and 30 deletions
View File
+99
View File
@@ -0,0 +1,99 @@
"""Tests for the CivitAI stats database layer."""
import pytest
from pathlib import Path
@pytest.fixture
def db_path(tmp_path):
return tmp_path / "test_stats.db"
@pytest.fixture
def stats_db(db_path):
from stats_db import StatsDB
db = StatsDB(db_path)
db.init()
return db
def test_init_creates_table(stats_db, db_path):
"""DB file is created and table exists."""
assert db_path.exists()
def test_upsert_and_get_single(stats_db):
"""Insert a stat row and retrieve it by sha256."""
stats_db.upsert("abc123", {
"civitai_model_id": 1,
"civitai_version_id": 10,
"download_count": 5000,
"rating": 4.8,
"rating_count": 120,
"thumbs_up_count": 89,
})
result = stats_db.get_by_hashes(["abc123"])
assert "abc123" in result
assert result["abc123"]["download_count"] == 5000
assert result["abc123"]["rating"] == 4.8
assert result["abc123"]["thumbs_up_count"] == 89
def test_upsert_updates_existing(stats_db):
"""Upserting same sha256 updates values."""
stats_db.upsert("abc123", {"download_count": 100, "rating": 3.0,
"rating_count": 10, "thumbs_up_count": 5})
stats_db.upsert("abc123", {"download_count": 200, "rating": 4.0,
"rating_count": 20, "thumbs_up_count": 15})
result = stats_db.get_by_hashes(["abc123"])
assert result["abc123"]["download_count"] == 200
def test_get_by_hashes_batch(stats_db):
"""Batch retrieval returns only matching hashes."""
for i in range(5):
stats_db.upsert(f"hash_{i}", {"download_count": i * 100,
"rating": 0, "rating_count": 0,
"thumbs_up_count": 0})
result = stats_db.get_by_hashes(["hash_1", "hash_3", "nonexistent"])
assert len(result) == 2
assert "hash_1" in result
assert "hash_3" in result
assert "nonexistent" not in result
def test_get_by_hashes_empty_input(stats_db):
"""Empty input returns empty dict."""
assert stats_db.get_by_hashes([]) == {}
def test_upsert_batch(stats_db):
"""Batch upsert inserts multiple rows."""
rows = [
("h1", {"civitai_model_id": 1, "download_count": 100,
"rating": 4.0, "rating_count": 10, "thumbs_up_count": 5}),
("h2", {"civitai_model_id": 2, "download_count": 200,
"rating": 3.5, "rating_count": 20, "thumbs_up_count": 15}),
]
stats_db.upsert_batch(rows)
result = stats_db.get_by_hashes(["h1", "h2"])
assert len(result) == 2
assert result["h1"]["download_count"] == 100
assert result["h2"]["download_count"] == 200
def test_get_all_stats(stats_db):
"""get_all returns every row keyed by sha256."""
stats_db.upsert("a", {"download_count": 1, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
stats_db.upsert("b", {"download_count": 2, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
result = stats_db.get_all()
assert len(result) == 2
def test_stats_count(stats_db):
"""count returns number of rows."""
assert stats_db.count() == 0
stats_db.upsert("a", {"download_count": 1, "rating": 0,
"rating_count": 0, "thumbs_up_count": 0})
assert stats_db.count() == 1