feat: add add_input_provenance function for test setups

provos · Dec 10, 2024 · bbb0f4b · bbb0f4b
1 parent 813d11b
commit bbb0f4b
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/src/planai/testing/test_helpers.py b/src/planai/testing/test_helpers.py
@@ -92,3 +92,8 @@ def assert_published_task_types(self, expected_types: List[Type[Task]]):
         assert (
             actual_types == expected_types
         ), f"Expected task types {expected_types}, got {actual_types}"
+
+
+def add_input_provenance(input_task: Task, provenance: Task) -> Task:
+    input_task._input_provenance.append(provenance)
+    return input_task