feat(trials): stop_trial and clear trials on reset
This commit is contained in:
@@ -69,3 +69,15 @@ def test_reset_empty_is_noop(tracker):
|
||||
tracker.start_trial("Pack")
|
||||
tracker.reset_trials_for(set())
|
||||
assert tracker.get_trials()[0]["unused_boot_days"] == 0
|
||||
|
||||
|
||||
def test_stop_trial_removes_row(tracker):
|
||||
tracker.start_trial("Pack")
|
||||
tracker.stop_trial("Pack")
|
||||
assert tracker.get_trials() == []
|
||||
|
||||
|
||||
def test_reset_clears_trials(tracker):
|
||||
tracker.start_trial("Pack")
|
||||
tracker.reset()
|
||||
assert tracker.get_trials() == []
|
||||
|
||||
+12
@@ -469,6 +469,17 @@ class UsageTracker:
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def stop_trial(self, package):
|
||||
"""End a trial (package became permanent or was disabled)."""
|
||||
with self._lock:
|
||||
self._ensure_db()
|
||||
conn = self._connect()
|
||||
try:
|
||||
conn.execute("DELETE FROM trial_packages WHERE package = ?", (package,))
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def reset(self):
|
||||
"""Clear all tracked data."""
|
||||
with self._lock:
|
||||
@@ -478,6 +489,7 @@ class UsageTracker:
|
||||
conn.execute("DELETE FROM node_usage")
|
||||
conn.execute("DELETE FROM prompt_log")
|
||||
conn.execute("DELETE FROM model_usage")
|
||||
conn.execute("DELETE FROM trial_packages")
|
||||
conn.commit()
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
Reference in New Issue
Block a user