#!/usr/bin/env python3 # -*- coding: utf-8 -*- """Tests for durable batch history and spend-safe ops.""" from __future__ import annotations import tempfile import unittest from pathlib import Path from types import SimpleNamespace from unittest import mock import util.translation as T import util.batch_history as BH class BatchHistoryTestBase(unittest.TestCase): """Isolate batch JSON files to a temp dir.""" def setUp(self): self._tmp = tempfile.TemporaryDirectory() tmp = Path(self._tmp.name) self._orig = { "QUEUE": T.BATCH_QUEUE_FILE, "STATE": T.BATCH_STATE_FILE, "RESULTS": T.BATCH_RESULTS_FILE, "LOCK": T.BATCH_LOCK_FILE, "HISTORY": BH.BATCH_HISTORY_FILE, "results_mem": T._batch_results, "pending": dict(T._batch_queue_pending), } T.BATCH_QUEUE_FILE = tmp / "batch_requests.json" T.BATCH_STATE_FILE = tmp / "batch_state.json" T.BATCH_RESULTS_FILE = tmp / "batch_results.json" T.BATCH_LOCK_FILE = tmp / "batch_files.lock" BH.BATCH_HISTORY_FILE = tmp / "batch_history.json" # batch_history imports queue/state/results paths at call time via T.* — # but it also imported BATCH_* as names. Rebind module-level aliases. BH.BATCH_QUEUE_FILE = T.BATCH_QUEUE_FILE BH.BATCH_STATE_FILE = T.BATCH_STATE_FILE BH.BATCH_RESULTS_FILE = T.BATCH_RESULTS_FILE T._batch_results = None T._batch_queue_pending = {} def tearDown(self): T.BATCH_QUEUE_FILE = self._orig["QUEUE"] T.BATCH_STATE_FILE = self._orig["STATE"] T.BATCH_RESULTS_FILE = self._orig["RESULTS"] T.BATCH_LOCK_FILE = self._orig["LOCK"] BH.BATCH_HISTORY_FILE = self._orig["HISTORY"] BH.BATCH_QUEUE_FILE = T.BATCH_QUEUE_FILE BH.BATCH_STATE_FILE = T.BATCH_STATE_FILE BH.BATCH_RESULTS_FILE = T.BATCH_RESULTS_FILE T._batch_results = self._orig["results_mem"] T._batch_queue_pending = self._orig["pending"] self._tmp.cleanup() class BatchRunStateTests(BatchHistoryTestBase): def test_none_when_empty(self): self.assertIsNone(T.batchRunState()) def test_queued_when_only_queue(self): T._write_batch_file(T.BATCH_QUEUE_FILE, {"k1": {"payload": "x", "language": "English", "params": {}}}) self.assertEqual(T.batchRunState(), "queued") def test_submitted_when_state_has_batches(self): T._write_batch_file( T.BATCH_STATE_FILE, {"batches": [{"id": "msgbatch_1", "custom_ids": {"req-000000": "k1"}}]}, ) self.assertEqual(T.batchRunState(), "submitted") def test_fetched_when_results_present(self): T._write_batch_file(T.BATCH_RESULTS_FILE, {"k1": {"text": "hi"}}) self.assertEqual(T.batchRunState(), "fetched") def test_fetched_when_state_status_fetched(self): T._write_batch_file( T.BATCH_STATE_FILE, {"status": "fetched", "batch_ids": ["msgbatch_1"], "batches": []}, ) self.assertEqual(T.batchRunState(), "fetched") class HistorySurvivalTests(BatchHistoryTestBase): def test_history_survives_fetch_marker_and_clear(self): custom_ids = {"req-000000": "cachekey1", "req-000001": "cachekey2"} BH.record_submit( [{"id": "msgbatch_abc", "custom_ids": custom_ids}], model="claude-sonnet-4-5", file_set=["Map001.json"], cost_estimate={"batch_cached_cost": 1.23, "model": "claude-sonnet-4-5"}, ) BH.record_fetch(["msgbatch_abc"], succeeded=2, errored=0, usage={"input_tokens": 10}, actual_cost=0.5) T._write_batch_file(T.BATCH_RESULTS_FILE, {"cachekey1": {"text": "A"}, "cachekey2": {"text": "B"}}) T._write_batch_file( T.BATCH_STATE_FILE, {"status": "fetched", "batch_ids": ["msgbatch_abc"], "batches": []}, ) T.clearBatchFiles() # Active files gone… self.assertFalse(T.BATCH_RESULTS_FILE.exists()) self.assertFalse(T.BATCH_STATE_FILE.exists()) # …but history retains custom_ids and is marked consumed. history = BH.read_history() entry = history["batches"][0] self.assertEqual(entry["id"], "msgbatch_abc") self.assertEqual(entry["custom_ids"], custom_ids) self.assertEqual(entry["status"], BH.STATUS_CONSUMED) self.assertEqual(entry["file_set"], ["Map001.json"]) def test_clear_does_not_wipe_history_file(self): BH.upsert_history_entry("msgbatch_keep", status=BH.STATUS_SUBMITTED, custom_ids={"a": "b"}) T.clearBatchFiles() self.assertTrue(BH.BATCH_HISTORY_FILE.exists()) self.assertEqual(len(BH.read_history()["batches"]), 1) class RedownloadTests(BatchHistoryTestBase): def test_redownload_rebuilds_results_from_custom_ids(self): custom_ids = {"req-000000": "keyA", "req-000001": "keyB"} BH.upsert_history_entry( "msgbatch_rd", status=BH.STATUS_ENDED, model="claude-sonnet-4-5", custom_ids=custom_ids, request_count=2, ) usage = SimpleNamespace( input_tokens=100, output_tokens=50, cache_read_input_tokens=10, cache_creation_input_tokens=20, thinking_tokens=5, ) msg = SimpleNamespace( content=[SimpleNamespace(text='{"Line1":"Hi"}')], usage=usage, ) ok_result = SimpleNamespace(type="succeeded", message=msg) row = SimpleNamespace(custom_id="req-000000", result=ok_result) row2_usage = SimpleNamespace( input_tokens=80, output_tokens=40, cache_read_input_tokens=0, cache_creation_input_tokens=0, thinking_tokens=0, ) msg2 = SimpleNamespace(content=[SimpleNamespace(text='{"Line1":"Yo"}')], usage=row2_usage) row2 = SimpleNamespace(custom_id="req-000001", result=SimpleNamespace(type="succeeded", message=msg2)) fake_batch = SimpleNamespace(id="msgbatch_rd", processing_status="ended") client = mock.MagicMock() client.messages.batches.retrieve.return_value = fake_batch client.messages.batches.results.return_value = [row, row2] with mock.patch.object(BH, "_get_anthropic_client", return_value=client): with mock.patch.object(BH, "getPricingConfig", return_value={"inputAPICost": 3.0, "outputAPICost": 15.0}): info = BH.redownload_batch("msgbatch_rd") self.assertEqual(info["succeeded"], 2) results = T._read_batch_file(T.BATCH_RESULTS_FILE) self.assertIn("keyA", results) self.assertIn("keyB", results) self.assertEqual(results["keyA"]["text"], '{"Line1":"Hi"}') state = T._read_batch_file(T.BATCH_STATE_FILE) self.assertEqual(state.get("status"), "fetched") self.assertEqual(T.batchRunState(), "fetched") entry = BH.read_history()["batches"][0] self.assertEqual(entry["status"], BH.STATUS_FETCHED) self.assertEqual(entry["custom_ids"], custom_ids) class CancelTests(BatchHistoryTestBase): def test_cancel_updates_history_and_active_state(self): T._write_batch_file( T.BATCH_STATE_FILE, { "batches": [ {"id": "msgbatch_c1", "custom_ids": {"req-000000": "k"}}, {"id": "msgbatch_c2", "custom_ids": {"req-000000": "k2"}}, ] }, ) BH.upsert_history_entry("msgbatch_c1", status=BH.STATUS_SUBMITTED, custom_ids={"req-000000": "k"}) BH.upsert_history_entry("msgbatch_c2", status=BH.STATUS_SUBMITTED, custom_ids={"req-000000": "k2"}) before = SimpleNamespace(id="msgbatch_c1", processing_status="in_progress") after = SimpleNamespace(id="msgbatch_c1", processing_status="canceling") client = mock.MagicMock() client.messages.batches.retrieve.return_value = before client.messages.batches.cancel.return_value = after with mock.patch.object(BH, "_get_anthropic_client", return_value=client): results = BH.cancel_batches(["msgbatch_c1"]) self.assertTrue(results[0]["ok"]) entry = next(e for e in BH.read_history()["batches"] if e["id"] == "msgbatch_c1") self.assertEqual(entry["status"], BH.STATUS_CANCELING) state = T._read_batch_file(T.BATCH_STATE_FILE) ids = [b["id"] for b in state.get("batches", [])] self.assertNotIn("msgbatch_c1", ids) self.assertIn("msgbatch_c2", ids) class UsageTests(BatchHistoryTestBase): def test_usage_sums_cache_and_thinking(self): BH.upsert_history_entry( "msgbatch_u", status=BH.STATUS_ENDED, model="claude-sonnet-4-5", custom_ids={"req-000000": "k"}, ) usage = SimpleNamespace( input_tokens=1000, output_tokens=200, cache_read_input_tokens=500, cache_creation_input_tokens=100, thinking_tokens=50, ) msg = SimpleNamespace(content=[SimpleNamespace(text="ok")], usage=usage) row = SimpleNamespace(custom_id="req-000000", result=SimpleNamespace(type="succeeded", message=msg)) client = mock.MagicMock() client.messages.batches.retrieve.return_value = SimpleNamespace(processing_status="ended") client.messages.batches.results.return_value = [row] with mock.patch.object(BH, "_get_anthropic_client", return_value=client): with mock.patch.object(BH, "getPricingConfig", return_value={"inputAPICost": 3.0, "outputAPICost": 15.0}): info = BH.usage_for_batch("msgbatch_u") u = info["usage"] self.assertEqual(u["input_tokens"], 1000) self.assertEqual(u["output_tokens"], 200) self.assertEqual(u["cache_read_input_tokens"], 500) self.assertEqual(u["cache_creation_input_tokens"], 100) self.assertEqual(u["thinking_tokens"], 50) self.assertIsInstance(info["actual_cost"], float) self.assertGreater(info["actual_cost"], 0) entry = BH.read_history()["batches"][0] self.assertEqual(entry["usage"]["thinking_tokens"], 50) class ActivateResumeTests(BatchHistoryTestBase): def test_activate_submitted_restores_state(self): custom_ids = {"req-000000": "k"} BH.upsert_history_entry( "msgbatch_act", status=BH.STATUS_SUBMITTED, custom_ids=custom_ids, model="claude-sonnet-4-5", file_set=["a.json"], ) state = BH.activate_for_resume("msgbatch_act") self.assertEqual(state, "submitted") disk = T._read_batch_file(T.BATCH_STATE_FILE) self.assertEqual(disk["batches"][0]["id"], "msgbatch_act") self.assertEqual(disk["batches"][0]["custom_ids"], custom_ids) if __name__ == "__main__": unittest.main()