diff --git a/tests/modules/test_module.py b/tests/modules/test_module.py index 3640a3e4..1515361b 100644 --- a/tests/modules/test_module.py +++ b/tests/modules/test_module.py @@ -14,6 +14,22 @@ def modify_value(_: str, record: dict, attr: str, value): record[attr] = value +def use_master_record( + _: str, record: dict, master_record: dict, target_attr: str, source_attr: str +): + """Hook that uses master record to copy a value from master to snapshot. + + Only applies when source attribute in master record has value starting with "master_" + to avoid interfering with other test cases. + """ + if source_attr in master_record: + # Get the value from master record + master_value = master_record[source_attr].get("v", None) + if master_value is not None and str(master_value).startswith("master_"): + # Append a suffix to demonstrate master record was used + record[target_attr] = f"{master_value}_from_master" + + dummy_hook_abc = update_wrapper(partial(modify_value, attr="data2", value="abc"), modify_value) dummy_hook_def = update_wrapper(partial(modify_value, attr="data1", value="def"), modify_value) @@ -67,3 +83,15 @@ def __init__( depends_on=[], may_change=[["data1"]], ) + + # Testing register_correlation_hook_with_master_record + # This hook should copy data1 from master record to data2 with a suffix + registrar.register_correlation_hook_with_master_record( + update_wrapper( + partial(use_master_record, target_attr="data2", source_attr="data1"), + use_master_record, + ), + "A", + depends_on=[], + may_change=[["data2"]], + ) diff --git a/tests/test_api/test_snapshots.py b/tests/test_api/test_snapshots.py index 8bdd4119..e9cd59f5 100644 --- a/tests/test_api/test_snapshots.py +++ b/tests/test_api/test_snapshots.py @@ -42,6 +42,9 @@ def make_dp(type, id, attr, v, time=False): make_dp("C", "c1", "ds", {"eid": "d1"}, time=True), make_dp("C", "c1", "data1", "inita"), make_dp("C", "c1", "data2", "inita"), + # For test_master_record_hook (A-423) + make_dp("A", 423, "data1", "master_test"), + make_dp("A", 423, "data2", "placeholder"), ] res = cls.push_datapoints(entity_datapoints) if res.status_code != 200: @@ -79,3 +82,17 @@ def test_hook_dependency_value_forwarding(self): for snapshot in data.snapshots: self.assertEqual(snapshot["data1"], "modifd") self.assertEqual(snapshot["data2"], "modifc") + + def test_master_record_hook(self): + """ + Test that hooks registered via register_correlation_hook_with_master_record + correctly receive the master record parameter. + """ + # Entity A-423 has data1="master_test" in its master record + # The master record hook should copy data1 from master and append "_from_master" + data = self.get_entity_data("entity/A/423", EntityEidData) + self.assertGreater(len(data.snapshots), 0) + for snapshot in data.snapshots: + self.assertEqual(snapshot["data1"], "master_test") + # The hook should have set data2 to data1 from master record + "_from_master" + self.assertEqual(snapshot["data2"], "master_test_from_master")