| """Unit tests for RecallTrace.""" |
|
|
| from __future__ import annotations |
|
|
| import unittest |
|
|
| from env.env import RecallTraceEnv |
| from grader.grader import evaluate_action_plan |
|
|
|
|
| class RecallTraceEnvTests(unittest.TestCase): |
| def test_phase1_plan_scores_high(self) -> None: |
| grade = evaluate_action_plan( |
| "phase1_direct_recall", |
| [ |
| {"type": "trace_lot", "lot_id": "LotA"}, |
| {"type": "inspect_node", "node_id": "warehouse"}, |
| {"type": "inspect_node", "node_id": "store1"}, |
| {"type": "inspect_node", "node_id": "store2"}, |
| {"type": "quarantine", "node_id": "warehouse", "lot_id": "LotA", "quantity": 100}, |
| {"type": "quarantine", "node_id": "store1", "lot_id": "LotA", "quantity": 50}, |
| {"type": "quarantine", "node_id": "store2", "lot_id": "LotA", "quantity": 20}, |
| {"type": "notify", "node_id": "all"}, |
| {"type": "finalize"}, |
| ], |
| ) |
| self.assertGreaterEqual(grade.score, 0.95) |
| self.assertTrue(grade.success) |
|
|
| def test_phase2_trace_reveals_relabels(self) -> None: |
| env = RecallTraceEnv(task_id="phase2_relabel_recall") |
| env.reset() |
| observation, reward, done, info = env.step({"type": "trace_lot", "lot_id": "LotA"}) |
| self.assertFalse(done) |
| self.assertGreater(reward, 0) |
| self.assertEqual(info["matched_lots"], ["LotA", "LotA_R1", "LotA_R2"]) |
| self.assertIn("store3", observation.trace_results["LotA"]["affected_nodes"]) |
|
|
| def test_phase3_mixed_inventory_requires_exact_quarantine(self) -> None: |
| env = RecallTraceEnv(task_id="phase3_mixed_shipments") |
| env.reset() |
| env.step({"type": "trace_lot", "lot_id": "LotA"}) |
| env.step({"type": "inspect_node", "node_id": "crossdock"}) |
| _, reward, _, info = env.step({"type": "quarantine", "node_id": "crossdock", "lot_id": "LotBlend", "quantity": 15}) |
| self.assertLess(reward, 0) |
| self.assertEqual(info["target_contaminated_quantity"], 12) |
|
|
| def test_phase3_full_plan_scores_high(self) -> None: |
| grade = evaluate_action_plan( |
| "phase3_mixed_shipments", |
| [ |
| {"type": "trace_lot", "lot_id": "LotA"}, |
| {"type": "inspect_node", "node_id": "warehouse"}, |
| {"type": "inspect_node", "node_id": "crossdock"}, |
| {"type": "inspect_node", "node_id": "store1"}, |
| {"type": "inspect_node", "node_id": "store2"}, |
| {"type": "inspect_node", "node_id": "store3"}, |
| {"type": "quarantine", "node_id": "warehouse", "lot_id": "LotA", "quantity": 30}, |
| {"type": "quarantine", "node_id": "crossdock", "lot_id": "LotBlend", "quantity": 12}, |
| {"type": "quarantine", "node_id": "store1", "lot_id": "LotA", "quantity": 10}, |
| {"type": "quarantine", "node_id": "store2", "lot_id": "LotBlend", "quantity": 8}, |
| {"type": "quarantine", "node_id": "store3", "lot_id": "LotBlend", "quantity": 4}, |
| {"type": "notify", "node_id": "all"}, |
| {"type": "finalize"}, |
| ], |
| ) |
| self.assertGreaterEqual(grade.score, 0.95) |
| self.assertTrue(grade.final_info["all_affected_stock_quarantined"]) |
|
|
|
|
| if __name__ == "__main__": |
| unittest.main() |
|
|
