Spaces:
Running
Running
| import importlib.util | |
| import json | |
| from pathlib import Path | |
| import tempfile | |
| import unittest | |
| from unittest import mock | |
| from types import SimpleNamespace | |
| spec = importlib.util.spec_from_file_location("study_tracking", Path(__file__).parents[1] / "tracking.py") | |
| tracking = importlib.util.module_from_spec(spec) | |
| spec.loader.exec_module(tracking) | |
| class FakeRemote: | |
| def __init__(self): | |
| self.fail = False | |
| self.visible = True | |
| self.logs = {} | |
| def predict(self, api_name, **kwargs): | |
| if self.fail: | |
| raise ConnectionError("secret must not appear in stored error") | |
| if api_name == "/bulk_log": | |
| for row in kwargs["logs"]: | |
| self.logs[row["log_id"]] = row | |
| return None | |
| return [{"value": row["metrics"]["tracking/event_sequence"]} | |
| for row in self.logs.values() if self.visible and row["run_id"] == kwargs["run_id"]] | |
| class TrackingTests(unittest.TestCase): | |
| def setUp(self): | |
| self.auth = mock.patch.dict("sys.modules", {"huggingface_hub": SimpleNamespace(get_token=lambda: None)}) | |
| self.auth.start() | |
| self.temp = tempfile.TemporaryDirectory() | |
| self.client = FakeRemote() | |
| self.tracker = tracking.Tracker(self.temp.name, client=self.client) | |
| self.event = {"timestamp": "2026-10-09T00:00:00Z", "run_id": "jq-single-r1", | |
| "task": "jqlang__jq.b33a763", "condition": "single", "kind": "progress", | |
| "metrics": {"api_input_tokens": 123, "score": 0.2}} | |
| def tearDown(self): | |
| self.tracker.close() | |
| self.temp.cleanup() | |
| self.auth.stop() | |
| def test_replay_is_idempotent(self): | |
| self.tracker.enqueue(self.event) | |
| self.tracker.enqueue(self.event) | |
| self.assertEqual(self.tracker.pending(), 1) | |
| self.assertTrue(self.tracker.flush()) | |
| self.assertEqual(self.tracker.pending(), 0) | |
| self.assertEqual(len(self.client.logs), 1) | |
| def test_failure_preserves_outbox_and_redacts_error(self): | |
| self.tracker.enqueue(self.event) | |
| self.client.fail = True | |
| self.assertFalse(self.tracker.flush()) | |
| self.assertEqual(self.tracker.pending(), 1) | |
| status = json.loads((Path(self.temp.name) / "status.json").read_text()) | |
| self.assertEqual(status["last_error"], "ConnectionError") | |
| self.tracker.client = self.client | |
| self.client.fail = False | |
| self.assertTrue(self.tracker.flush()) | |
| def test_http_success_without_readback_is_not_acknowledged(self): | |
| self.tracker.enqueue(self.event) | |
| self.client.visible = False | |
| self.assertFalse(self.tracker.flush()) | |
| self.assertEqual(self.tracker.pending(), 1) | |
| def test_partial_line_waits_and_restart_keeps_cursor(self): | |
| path = Path(self.temp.name) / "input.jsonl" | |
| line = json.dumps(self.event) | |
| path.write_text(line) | |
| self.tracker.ingest(path) | |
| self.assertEqual(self.tracker.pending(), 0) | |
| path.write_text(line + "\n") | |
| self.tracker.ingest(path) | |
| self.tracker.ingest(path) | |
| self.assertEqual(self.tracker.pending(), 1) | |
| def test_raw_text_and_credentials_are_never_published(self): | |
| event = self.event | {"message": "secret", "config": {"api_key": "secret", "model": "Beam-501B-A23B", | |
| "prompt": "private prompt", "max_tokens": 8}} | |
| payload = tracking.sanitize(event) | |
| self.assertNotIn("secret", json.dumps(payload)) | |
| self.assertNotIn("private prompt", json.dumps(payload)) | |
| self.assertEqual(payload["config"]["max_tokens"], 8) | |
| if __name__ == "__main__": | |
| unittest.main() | |