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