Reduce nesting in _get_single_mapping

Fix pytest.raises logic error in test_ioc.py
DiamondLightSource · Sep 17, 2024 · 00fbd1d · 00fbd1d
1 parent d079439
commit 00fbd1d
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 16 deletions.
diff --git a/src/fastcs/mapping.py b/src/fastcs/mapping.py
@@ -38,25 +38,21 @@ def _walk_mappings(controller: BaseController) -> Iterator[SingleMapping]:
 
 
 def _get_single_mapping(controller: BaseController) -> SingleMapping:
-    scan_methods = {}
-    put_methods = {}
-    command_methods = {}
-    attributes = {}
+    scan_methods: dict[str, Scan] = {}
+    put_methods: dict[str, Put] = {}
+    command_methods: dict[str, Command] = {}
+    attributes: dict[str, Attribute] = {}
     for attr_name in dir(controller):
         attr = getattr(controller, attr_name)
         match attr:
-            case WrappedMethod(fastcs_method=fastcs_method):
-                if fastcs_method.enabled:
-                    match fastcs_method:
-                        case Put():
-                            put_methods[attr_name] = fastcs_method
-                        case Scan():
-                            scan_methods[attr_name] = fastcs_method
-                        case Command():
-                            command_methods[attr_name] = fastcs_method
-            case Attribute():
-                if attr.enabled:
-                    attributes[attr_name] = attr
+            case WrappedMethod(fastcs_method=Put(enabled=True) as put_method):
+                put_methods[attr_name] = put_method
+            case WrappedMethod(fastcs_method=Scan(enabled=True) as scan_method):
+                scan_methods[attr_name] = scan_method
+            case WrappedMethod(fastcs_method=Command(enabled=True) as command_method):
+                command_methods[attr_name] = command_method
+            case Attribute(enabled=True):
+                attributes[attr_name] = attr
 
     return SingleMapping(
         controller, scan_methods, put_methods, command_methods, attributes

diff --git a/tests/backends/epics/test_ioc.py b/tests/backends/epics/test_ioc.py
@@ -201,4 +201,5 @@ def test_long_pv_names_discarded(mocker: MockerFixture):
             always_update=True,
             on_update=mocker.ANY,
         )
+    with pytest.raises(AssertionError):
         builder.longIn.assert_called_once_with(f"{DEVICE}:{long_rw_pv_name}_RBV")