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.start_trial("Pack")
|
||||||
tracker.reset_trials_for(set())
|
tracker.reset_trials_for(set())
|
||||||
assert tracker.get_trials()[0]["unused_boot_days"] == 0
|
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:
|
finally:
|
||||||
conn.close()
|
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):
|
def reset(self):
|
||||||
"""Clear all tracked data."""
|
"""Clear all tracked data."""
|
||||||
with self._lock:
|
with self._lock:
|
||||||
@@ -478,6 +489,7 @@ class UsageTracker:
|
|||||||
conn.execute("DELETE FROM node_usage")
|
conn.execute("DELETE FROM node_usage")
|
||||||
conn.execute("DELETE FROM prompt_log")
|
conn.execute("DELETE FROM prompt_log")
|
||||||
conn.execute("DELETE FROM model_usage")
|
conn.execute("DELETE FROM model_usage")
|
||||||
|
conn.execute("DELETE FROM trial_packages")
|
||||||
conn.commit()
|
conn.commit()
|
||||||
finally:
|
finally:
|
||||||
conn.close()
|
conn.close()
|
||||||
|
|||||||
Reference in New Issue
Block a user