test: add tier coverage for safe_to_remove and consider_removing

This commit is contained in:
2026-04-09 17:58:26 +02:00
parent 800692d190
commit ff42016009
+29
View File
@@ -1,6 +1,8 @@
import pytest import pytest
import tempfile import tempfile
import os import os
from unittest.mock import patch
from datetime import datetime, timezone, timedelta
from tracker import UsageTracker from tracker import UsageTracker
@@ -72,3 +74,30 @@ def test_get_model_stats_sorted_by_status(tracker):
statuses = [m["status"] for m in models] statuses = [m["status"] for m in models]
# unused_new (2) comes before used (3) in STATUS_ORDER # unused_new (2) comes before used (3) in STATUS_ORDER
assert statuses.index("unused_new") < statuses.index("used") assert statuses.index("unused_new") < statuses.index("used")
def test_get_model_stats_safe_to_remove(tracker):
"""A model last used 70 days ago should be classified safe_to_remove."""
tracker.record_model_usage([("old.safetensors", "checkpoints")])
installed = {"checkpoints": ["old.safetensors"]}
# Patch datetime in tracker so "now" is 70 days after last_seen
future_now = datetime.now(timezone.utc) + timedelta(days=70)
with patch("tracker.datetime") as mock_dt:
mock_dt.now.return_value = future_now
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "safe_to_remove"
def test_get_model_stats_consider_removing(tracker):
"""A model last used 40 days ago should be classified consider_removing."""
tracker.record_model_usage([("medium.safetensors", "checkpoints")])
installed = {"checkpoints": ["medium.safetensors"]}
future_now = datetime.now(timezone.utc) + timedelta(days=40)
with patch("tracker.datetime") as mock_dt:
mock_dt.now.return_value = future_now
result = tracker.get_model_stats(installed)
assert result[0]["models"][0]["status"] == "consider_removing"