MNT/TST: adjust Client, tests to expect EpicsData

pcdshub · Aug 5, 2024 · 416c645 · 416c645
1 parent e081276
commit 416c645
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 10 deletions.
diff --git a/superscore/client.py b/superscore/client.py
@@ -8,12 +8,11 @@
 
 from superscore.backends import get_backend
 from superscore.backends.core import _Backend
-from superscore.control_layers import ControlLayer
+from superscore.control_layers import ControlLayer, EpicsData
 from superscore.control_layers.status import TaskStatus
 from superscore.errors import CommunicationError
 from superscore.model import (Collection, Entry, Nestable, Parameter, Readback,
                               Setpoint, Snapshot)
-from superscore.type_hints import AnyEpicsType
 from superscore.utils import build_abs_path
 
 logger = logging.getLogger(__name__)
@@ -290,7 +289,7 @@ def _gather_data(
     def _build_snapshot(
         self,
         coll: Collection,
-        values: Dict[str, AnyEpicsType],
+        values: Dict[str, EpicsData],
     ) -> Snapshot:
         """
         Traverse a Collection, assembling a Snapshot using pre-fetched data
@@ -300,7 +299,7 @@ def _build_snapshot(
         ----------
         coll : Collection
             The collection being saved
-        values : Dict[str, AnyEpicsType]
+        values : Dict[str, EpicsData]
             A dictionary mapping PV names to pre-fetched values
 
         Returns
@@ -319,17 +318,23 @@ def _build_snapshot(
                 child = self.backend.get(child)
             if isinstance(child, Parameter):
                 if child.readback is not None:
+                    edata = values[child.readback.pv_name] or EpicsData(data=None)
                     readback = Readback(
                         pv_name=child.readback.pv_name,
                         description=child.readback.description,
-                        data=values[child.readback.pv_name]
+                        data=edata.data,
+                        status=edata.status,
+                        severity=edata.severity
                     )
                 else:
                     readback = None
+                edata = values[child.pv_name] or EpicsData(data=None)
                 setpoint = Setpoint(
                     pv_name=child.pv_name,
                     description=child.description,
-                    data=values[child.pv_name],
+                    data=edata.data,
+                    status=edata.status,
+                    severity=edata.severity,
                     readback=readback
                 )
                 snapshot.children.append(setpoint)
@@ -338,9 +343,12 @@ def _build_snapshot(
 
         snapshot.meta_pvs = []
         for pv in Collection.meta_pvs:
+            edata = values[pv] or EpicsData(data=None)
             readback = Readback(
-                pv_name=readback.pv_name,
-                data=values[readback.pv_name]
+                pv_name=pv,
+                data=edata.data,
+                status=edata.status,
+                severity=edata.severity,
             )
             snapshot.meta_pvs.append(readback)
 

diff --git a/superscore/tests/test_client.py b/superscore/tests/test_client.py
@@ -6,6 +6,7 @@
 
 from superscore.backends.filestore import FilestoreBackend
 from superscore.client import Client
+from superscore.control_layers import EpicsData
 from superscore.errors import CommunicationError
 from superscore.model import Parameter, Readback, Root, Setpoint
 
@@ -84,7 +85,7 @@ def test_snap(
     coll = sample_database.entries[2]
     coll.children.append(parameter_with_readback)
 
-    get_mock.side_effect = range(5)
+    get_mock.side_effect = [EpicsData(i) for i in range(5)]
     snapshot = mock_client.snap(coll)
     assert get_mock.call_count == 5
     assert all([snapshot.children[i].data == i for i in range(4)])  # children saved in order
@@ -97,7 +98,8 @@ def test_snap(
 @patch('superscore.control_layers.core.ControlLayer._get_one')
 def test_snap_exception(get_mock, mock_client: Client, sample_database: Root):
     coll = sample_database.entries[2]
-    get_mock.side_effect = [0, 1, CommunicationError, 3, 4]
+    get_mock.side_effect = [EpicsData(0), EpicsData(1), CommunicationError,
+                            EpicsData(3), EpicsData(4)]
     snapshot = mock_client.snap(coll)
     assert snapshot.children[2].data is None