import copy from collections import Counter import unittest from unittest.mock import patch import seed_registry import trial_plan as m class TrialPlan(unittest.TestCase): @classmethod def setUpClass(cls): cls.original = m.build_trial_plan(461) def test_complete_reconstruction_and_budget(self): self.assertTrue(m.validate_trial_plan(self.original)) budget = self.original["budget"] self.assertEqual(budget["training_control_reversal_epochs"], 68000) self.assertEqual(budget["training_control_reversal_steps"], 99600000) self.assertEqual(budget["minimum_biological_seconds_excluding_evaluation"], 10760) self.assertEqual(budget["minimum_steps_excluding_evaluation"], 107600000) self.assertFalse(budget["evaluation_probe_inventory_complete"]) self.assertIsNone(budget["evaluation_epochs"]) self.assertFalse(self.original["plan"]["training_executed"]) def test_exact_cells_five_training_replicates_not_evaluation_streams(self): cells = {} for block in self.original["plan"]["blocks"]: key = block["mode"], block["dan_type"], block["initial_reinforced_cue"] cells.setdefault(key, []).append(block["replicate"]) self.assertEqual(len(block["evaluation_streams"]), 2) self.assertEqual(len(cells), 8) self.assertTrue(all(replicates == [1, 2, 3, 4, 5] for replicates in cells.values())) def test_all_epoch_control_roles_and_matched_delayed_slots(self): for block in self.original["plan"]["blocks"]: for suite in block["suites"]: assay = suite["assay"] arms = {a["arm"]: a for a in suite["arms"]} self.assertEqual(tuple(arms), m.REVERSAL_ARMS if assay == "reversal" else m.ARMS) expected_steps, expected_us = (1800, 1500) if assay == "delayed" else (1200, 900) for arm, spec in arms.items(): epochs = spec["epochs"] self.assertEqual(len(epochs), 150 if arm == "unpaired" else 100) cues = Counter(e["cs"]["cue"] for e in epochs if e["cs"] is not None) self.assertEqual(cues, Counter() if arm == "dan_only" else Counter(A=50, B=50)) us = [e for e in epochs if e["us"]["requested"]] self.assertEqual(len(us), 0 if arm == "cs_only" else 50) self.assertEqual(len({e["us"]["inventory_id"] for e in us}), len(us)) for i, e in enumerate(epochs): self.assertEqual(e["epoch_index"], i) self.assertEqual(e["steps"], expected_steps) self.assertEqual(e["us"]["nominal_slot_tick"], expected_us) self.assertEqual(e["us"]["tick"], expected_us if e["us"]["requested"] else None) self.assertIsNone(e["us"]["rng_seed"]) self.assertEqual(e["us"]["rng_draws"], 0) self.assertEqual(e["update_enabled"], arm != "frozen") self.assertEqual(e["D_command"], 1.0 if e["us"]["requested"] and arm != "d_zero" else 0.0) if arm == "random_reward": self.assertEqual(Counter(e["cs"]["cue"] for e in us), Counter(A=25, B=25)) if arm == "unpaired": self.assertTrue(all(e["cs"] is None for e in us)) if arm == "dan_only": self.assertEqual(sum(not e["us"]["requested"] for e in epochs), 50) def test_alias_binding_same_occurrence_even_unpaired_insertion(self): for block in self.original["plan"]["blocks"]: for suite in block["suites"]: base = suite["arms"][0] lookup = {(e["cs"]["cue"], e["cs"]["occurrence"]): e["cs"] for e in base["epochs"]} base_order = [(e["cs"]["cue"], e["cs"]["occurrence"]) for e in base["epochs"]] for arm in suite["arms"][1:]: actual_order = [] for e in arm["epochs"]: if e["cs"] is None: continue cs = e["cs"] source = lookup[(cs["cue"], cs["occurrence"])] self.assertEqual(cs["seed"], source["seed"]) self.assertEqual(cs["source_stream_id"], source["stream_id"]) self.assertNotEqual(cs["stream_id"], source["stream_id"]) actual_order.append((cs["cue"], cs["occurrence"])) if arm["arm"] != "dan_only": self.assertEqual(actual_order, base_order) def test_reversal_not_extra_acquisition_or_new_compartment(self): for b in self.original["plan"]["blocks"]: reversal, continuation = b["suites"][2]["arms"] original = b["initial_reinforced_cue"] other = b["initial_neutral_cue"] for spec, cue in ((reversal, other), (continuation, original)): self.assertEqual(spec["delta_start"], {"kind": "fork", "source": b["block_id"] + "/normal/paired/terminal-delta"}) self.assertEqual({e["cs"]["cue"] for e in spec["epochs"] if e["us"]["requested"]}, {cue}) def test_split_and_replicate_seeds_are_independent_fullwidth(self): manifest = self.original["seed_manifest"] all_seeds = [r["seed"] for r in manifest["independent_streams"]] self.assertEqual(len(all_seeds), len(set(all_seeds))) by_block = {} by_split = {} for row in manifest["independent_streams"]: by_block.setdefault(row["key"][0], set()).add(row["seed"]) by_split.setdefault(row["split"], set()).add(row["seed"]) self.assertEqual(len(by_block), 40) blocks = list(by_block) for i, block in enumerate(blocks): for other in blocks[i + 1:]: self.assertFalse(by_block[block] & by_block[other]) self.assertFalse(by_split["training"] & by_split["evaluation"]) self.assertTrue(any(seed > 2**32 for seed in all_seeds)) for block in self.original["plan"]["blocks"]: for probe in block["evaluation_streams"]: self.assertIn(probe["seed"], by_split["evaluation"]) self.assertEqual(probe["seed_bits"], 63) def test_generic_valid_alias_can_still_be_protocol_illegal(self): for source, destination in ( (("b-rep1", "normal", "CS", "A", "00", "paired"), ("b-rep2", "normal", "CS", "A", "00", "frozen")), (("b", "normal", "US", "A", "00", "paired"), ("b", "normal", "CS", "A", "00", "frozen")), (("b", "normal", "CS", "A", "00", "paired"), ("b", "delayed", "CS", "A", "00", "frozen")), (("b", "normal", "CS", "A", "00", "paired"), ("b", "normal", "CS", "B", "00", "frozen")), (("b", "normal", "CS", "A", "00", "paired"), ("b", "normal", "CS", "A", "01", "frozen")), (("b", "normal", "CS", "A", "00", "paired"), ("b", "normal", "CS", "A", "00", "undeclared-control")), ): r = seed_registry.SeedRegistry(1) r.add("training", source) r.alias("training", destination, "training", source, reason="Synthetic prohibited crossing") self.assertTrue(seed_registry.validate_manifest(r.manifest())) with self.assertRaises(ValueError): m._validate_alias_permissions(r.manifest()) def test_mutations_fail_even_after_forging_payload_sha(self): def first(p): return p["plan"]["blocks"][0]["suites"][0]["arms"][0]["epochs"][0] mutations = ( lambda p: p["plan"]["blocks"].pop(), lambda p: p["plan"]["blocks"].__setitem__(1, copy.deepcopy(p["plan"]["blocks"][0])), lambda p: p["plan"]["blocks"][0]["suites"][0]["arms"][0]["epochs"].pop(), lambda p: first(p)["cs"].update(seed=first(p)["cs"]["seed"] & 0xffffffff), lambda p: first(p)["cs"].update(seed_bits=32), lambda p: first(p)["cs"].update(odor="post-outcome-selected odor"), lambda p: p["plan"].update(execution_authorized=0), lambda p: p["budget"].update(evaluation_epochs=0), ) for index, mutate in enumerate(mutations): payload = copy.deepcopy(self.original) mutate(payload) payload["payload_sha256"] = m._sha({k: v for k, v in payload.items() if k != "payload_sha256"}) with self.subTest(mutation=index), self.assertRaises(ValueError): m.validate_trial_plan(payload) def test_deterministic_small_shuffle_and_fullwidth_seed(self): a = m._Shuffle(23).shuffled(range(30)) self.assertEqual(a, m._Shuffle(23).shuffled(range(30))) self.assertNotEqual(a, m._Shuffle(23 + 2**32).shuffled(range(30))) self.assertEqual(sorted(a), list(range(30))) with self.assertRaises(ValueError): m._Shuffle(True) def test_integer_rejection_advances_counter_not_seed(self): shuffle = m._Shuffle(23 + 2**32) with patch.object(m, "_sha", side_effect=["f" * 64, "0" * 64]) as hashed: self.assertEqual(shuffle.below(3), 0) self.assertEqual(shuffle.counter, 2) self.assertEqual(hashed.call_args_list[0].args[0], [m.SHUFFLE_VERSION, 23 + 2**32, 0]) self.assertEqual(hashed.call_args_list[1].args[0], [m.SHUFFLE_VERSION, 23 + 2**32, 1]) if __name__ == "__main__": unittest.main()