Refactor tests and add tests for the beamline name mapping

src/dodal/beamlines/__init__.py

@@ -22,12 +22,19 @@ def all_beamline_modules() -> Iterable[str]:
     Returns:
         Iterable[str]: Generator of beamline module names
     """
+
+    # This is done bty inspecting file names rather than modules to avoid
+    # premature importing
     spec = importlib.util.find_spec(__name__)
     if spec is not None:
         search_paths = [Path(path) for path in spec.submodule_search_locations]
         for path in search_paths:
             for subpath in path.glob("**/*"):
-                if subpath.name.endswith(".py") and subpath.name != "__init__.py":
+                if (
+                    subpath.name.endswith(".py")
+                    and subpath.name != "__init__.py"
+                    and ("__pycache__" not in str(subpath))
+                ):
                     yield subpath.with_suffix("").name
     else:
         raise KeyError(f"Unable to find {__name__} module")
@@ -42,7 +49,7 @@ def all_beamline_names() -> Iterable[str]:
     """
     inverse_mapping = _module_name_overrides()
     for module_name in all_beamline_modules():
-        yield from inverse_mapping.get(module_name, {module_name})
+        yield from inverse_mapping.get(module_name, set()).union({module_name})
 
 
 @lru_cache

tests/beamlines/unit_tests/test_mapping.py

@@ -0,0 +1,35 @@
+import pytest
+
+from dodal.beamlines import (
+    all_beamline_names,
+    module_name_for_beamline,
+)
+
+
+@pytest.mark.parametrize(
+    "beamline,expected_module",
+    {
+        "i03": "i03",
+        "s03": "i03",
+        "i04-1": "i04_1",
+        "i22": "i22",
+    }.items(),
+)
+def test_beamline_name_mapping(beamline: str, expected_module: str):
+    assert module_name_for_beamline(beamline) == expected_module
+
+
+def test_all_beamline_names_includes_non_overridden_modules():
+    beamlines = set(all_beamline_names())
+    assert "i22" in beamlines
+
+
+def test_all_beamline_names_includes_overriden_modules():
+    beamlines = set(all_beamline_names())
+    assert "i04-1" in beamlines
+
+
+def test_all_beamline_names_includes_overriden_default_modules():
+    beamlines = set(all_beamline_names())
+    assert "i03" in beamlines
+    assert "s03" in beamlines

tests/unit_tests/test_cli.py

@@ -1,12 +1,16 @@
+import os
 from unittest.mock import patch
 
 import pytest
-from click.testing import CliRunner
+from click.testing import CliRunner, Result
 
 from dodal import __version__
-from dodal.beamlines import all_beamline_names
 from dodal.cli import main
 
+# Test with an example beamline, device instantiation is already tested
+# in beamline unit tests
+EXAMPLE_BEAMLINE = "i22"
+
 
 @pytest.fixture
 def runner():
@@ -23,33 +27,30 @@ def test_cli_version(runner: CliRunner):
     assert result.stdout == f"{__version__}\n"
 
 
-def test_cli_connect(runner: CliRunner):
-    beamline = next(iter(all_beamline_names()))  # Test with an example beamline,
-    # device instantiation is already tested in beamline unit tests
+def test_cli_sets_beamline_environment_variable(runner: CliRunner):
+    with patch.dict(os.environ, clear=True):
+        _mock_connect(EXAMPLE_BEAMLINE, runner=runner)
+        assert os.environ["BEAMLINE"] == EXAMPLE_BEAMLINE
 
-    with patch(
-        "dodal.cli.make_all_devices",
-        return_value={f"device_{i}": object() for i in range(3)},
-    ):
-        result = runner.invoke(
-            main,
-            ["connect", beamline],
-            catch_exceptions=False,
-        )
+
+def test_cli_connect(runner: CliRunner):
+    result = _mock_connect(EXAMPLE_BEAMLINE, runner=runner)
     assert result.stdout.startswith("3 devices connected")
 
 
 def test_cli_connect_in_sim_mode(runner: CliRunner):
-    beamline = next(iter(all_beamline_names()))  # Test with an example beamline,
-    # device instantiation is already tested in beamline unit tests
+    result = _mock_connect("-s", EXAMPLE_BEAMLINE, runner=runner)
+    assert result.stdout.startswith("3 devices connected (sim mode)")
+
 
+def _mock_connect(*args, runner: CliRunner) -> Result:
     with patch(
         "dodal.cli.make_all_devices",
         return_value={f"device_{i}": object() for i in range(3)},
     ):
         result = runner.invoke(
             main,
-            ["connect", "-s", beamline],
+            ["connect"] + list(args),
             catch_exceptions=False,
         )
-    assert result.stdout.startswith("3 devices connected (sim mode)")
+    return result