Fix test suite to correctly regenerate examples across tests

src/DataSchemas/aind_behavior_services/launcher.py

@@ -176,8 +176,10 @@ def _print_header(self) -> None:
             f"TaskLogic ({self.task_logic_schema.__name__}) \
                 Schema Version: {self.task_logic_schema.model_construct().version}"
         )
-        print(f"Rig ({self.rig_schema.__name__}) \
-              Schema Version: {self.rig_schema.model_construct().version}")
+        print(
+            f"Rig ({self.rig_schema.__name__}) \
+              Schema Version: {self.rig_schema.model_construct().version}"
+        )
         print(
             f"Session ({self.session_schema.__name__}) \
                 Schema Version: {self.session_schema.model_construct().version}"
@@ -612,4 +614,4 @@ def _get_default_arg_parser() -> argparse.ArgumentParser:
             default=False,
         )
 
-        return parser
+        return parser

tests/__init__.py

@@ -0,0 +1,21 @@
+import glob
+import importlib.util
+from pathlib import Path
+from types import ModuleType
+
+EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
+
+
+def build_example(script_path: str) -> ModuleType:
+    module_name = Path(script_path).stem
+    spec = importlib.util.spec_from_file_location(module_name, script_path)
+    if spec is None:
+        raise ImportError(f"Can't find {script_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
+
+
+def build_examples(examples_dir: Path = EXAMPLES_DIR):
+    for script_path in glob.glob(str(examples_dir / "*.py")):
+        _ = build_example(script_path)

tests/test_bonsai.py

@@ -3,9 +3,10 @@
 from pathlib import Path
 from typing import Dict
 
-import test_examples
 from aind_behavior_services.utils import run_bonsai_process
 
+from . import build_examples
+
 
 class BonsaiTests(unittest.TestCase):
 
@@ -18,6 +19,7 @@ def test_deserialization(self):
             "water_valve",
         ]
 
+        build_examples()
         JSON_ROOT = Path("./local").resolve()
 
         workflow_props: Dict[str, str] = {}

tests/test_examples.py

@@ -2,24 +2,17 @@
 
 import glob
 import unittest
-from pathlib import Path
 
-EXAMPLES_DIR = Path(__file__).parents[1] / "examples"
+from . import EXAMPLES_DIR, build_example
 
 
 class ExampleTests(unittest.TestCase):
     """tests for examples"""
 
     def test_examples(self):
-        import importlib.util
-
         for script_path in glob.glob(str(EXAMPLES_DIR / "*.py")):
             with self.subTest(script_path=script_path):
-                module_name = Path(script_path).stem
-                spec = importlib.util.spec_from_file_location(module_name, script_path)
-                module = importlib.util.module_from_spec(spec)
-                spec.loader.exec_module(module)
-
+                module = build_example(script_path)
                 # Check if the module executed successfully
                 if hasattr(module, "__name__") and module.__name__ == "__main__":
                     self.assertEqual(module.__name__, "__main__", f"Script {script_path} failed to execute")

tests/test_water_valve.py

@@ -19,8 +19,7 @@ def test_calibration(self):
         _delta_times = [0.1, 0.2, 0.3, 0.4, 0.5]
         _slope = 10.1
         _offset = -0.3
-        _linear_model = lambda time: _slope * time + _offset
-        _water_weights = [_linear_model(x) for x in _delta_times]
+        _water_weights = [water_mock_model(x, _slope, _offset) for x in _delta_times]
         _inputs = [
             Measurement(valve_open_interval=0.5, valve_open_time=t[0], water_weight=[t[1]], repeat_count=1)
             for t in zip(_delta_times, _water_weights)
@@ -52,8 +51,7 @@ def test_calibration_on_null_output(self):
         _delta_times = [0.1, 0.2, 0.3, 0.4, 0.5]
         _slope = 10.1
         _offset = -0.3
-        _linear_model = lambda time: _slope * time + _offset
-        _water_weights = [_linear_model(x) for x in _delta_times]
+        _water_weights = [water_mock_model(x, _slope, _offset) for x in _delta_times]
         _inputs = WaterValveCalibrationInput(
             measurements=[
                 Measurement(valve_open_interval=0.5, valve_open_time=t[0], water_weight=[t[1]], repeat_count=1)
@@ -73,5 +71,9 @@ def test_calibration_on_null_output(self):
         self.assertAlmostEqual(1.0, calibration.output.r2, 2, "R2 is not almost equal")
 
 
+def water_mock_model(time: float, slope: float, offset: float) -> float:
+    return slope * time + offset
+
+
 if __name__ == "__main__":
     unittest.main()