feat: add model_usage table and tracker methods
Adds model_usage schema, record_model_usage(), get_raw_model_stats(), and get_model_stats() to UsageTracker, with reset() support and 8 passing tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
@@ -0,0 +1,74 @@
|
||||
import pytest
|
||||
import tempfile
|
||||
import os
|
||||
from tracker import UsageTracker
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def tracker(tmp_path):
|
||||
return UsageTracker(db_path=str(tmp_path / "test.db"))
|
||||
|
||||
|
||||
def test_record_and_retrieve_model_usage(tracker):
|
||||
tracker.record_model_usage([("dreamshaper.safetensors", "checkpoints")])
|
||||
tracker.record_model_usage([("dreamshaper.safetensors", "checkpoints")])
|
||||
|
||||
raw = tracker.get_raw_model_stats()
|
||||
assert len(raw) == 1
|
||||
assert raw[0]["model_name"] == "dreamshaper.safetensors"
|
||||
assert raw[0]["model_type"] == "checkpoints"
|
||||
assert raw[0]["count"] == 2
|
||||
|
||||
|
||||
def test_record_multiple_models(tracker):
|
||||
tracker.record_model_usage([
|
||||
("dreamshaper.safetensors", "checkpoints"),
|
||||
("vae.safetensors", "vae"),
|
||||
])
|
||||
raw = tracker.get_raw_model_stats()
|
||||
assert len(raw) == 2
|
||||
|
||||
|
||||
def test_reset_clears_model_usage(tracker):
|
||||
tracker.record_model_usage([("model.safetensors", "checkpoints")])
|
||||
tracker.reset()
|
||||
assert tracker.get_raw_model_stats() == []
|
||||
|
||||
|
||||
def test_empty_models_returns_empty(tracker):
|
||||
assert tracker.get_raw_model_stats() == []
|
||||
|
||||
|
||||
def test_get_model_stats_used(tracker):
|
||||
tracker.record_model_usage([("model.safetensors", "checkpoints")])
|
||||
installed = {"checkpoints": ["model.safetensors"]}
|
||||
result = tracker.get_model_stats(installed)
|
||||
assert len(result) == 1
|
||||
assert result[0]["model_type"] == "checkpoints"
|
||||
assert result[0]["models"][0]["status"] == "used"
|
||||
assert result[0]["models"][0]["count"] == 1
|
||||
|
||||
|
||||
def test_get_model_stats_never_used_new(tracker):
|
||||
installed = {"checkpoints": ["unused.safetensors"]}
|
||||
result = tracker.get_model_stats(installed)
|
||||
assert result[0]["models"][0]["status"] == "unused_new"
|
||||
assert result[0]["models"][0]["count"] == 0
|
||||
|
||||
|
||||
def test_get_model_stats_uninstalled(tracker):
|
||||
tracker.record_model_usage([("gone.safetensors", "checkpoints")])
|
||||
installed = {} # no longer on disk
|
||||
result = tracker.get_model_stats(installed)
|
||||
assert result[0]["models"][0]["status"] == "uninstalled"
|
||||
assert result[0]["models"][0]["installed"] is False
|
||||
|
||||
|
||||
def test_get_model_stats_sorted_by_status(tracker):
|
||||
tracker.record_model_usage([("active.safetensors", "checkpoints")])
|
||||
installed = {"checkpoints": ["active.safetensors", "unused.safetensors"]}
|
||||
result = tracker.get_model_stats(installed)
|
||||
models = result[0]["models"]
|
||||
statuses = [m["status"] for m in models]
|
||||
# unused_new (2) comes before used (3) in STATUS_ORDER
|
||||
assert statuses.index("unused_new") < statuses.index("used")
|
||||
Reference in New Issue
Block a user