Ethanfel
44088649d8
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>
100 lines
3.3 KiB
Python
100 lines
3.3 KiB
Python
"""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
|