fix(device_client): fix the on_message error handling

CederGroupHub · Aug 2, 2024 · 649ef05 · 649ef05
1 parent 4aec1b5
commit 649ef05
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/alab_management/device_manager.py b/alab_management/device_manager.py
@@ -372,10 +372,14 @@ def on_message(
         _body: bytes,
     ):
         """Callback function to handle a returned message from Device Manager."""
-        body = dill.loads(_body)
-
         f = self._waiting.pop(ObjectId(properties.correlation_id))
-        if body["status"] == "success":
-            f.set_result(body["result"])
-        else:
-            f.set_exception(body["result"])
+
+        try:
+            body = dill.loads(_body)
+
+            if body["status"] == "success":
+                f.set_result(body["result"])
+            else:
+                f.set_exception(body["result"])
+        except Exception as e:
+            f.set_exception(e)
diff --git a/tests/test_task_actor.py b/tests/test_task_actor.py
@@ -69,6 +69,7 @@ def compose_exp(exp_name):
                         "prev_tasks": [],
                         "parameters": {
                             "dest": "furnace_temp",
+                            "priority": 100,
                         },
                         "samples": ["test_sample"],
                     },