fix tests

radical-cybertools · Sep 22, 2023 · dd8420e · dd8420e
1 parent 48d2df4
commit dd8420e
Show file tree

Hide file tree

Showing 3 changed files with 7 additions and 8 deletions.
diff --git a/src/radical/pilot/agent/agent_0.py b/src/radical/pilot/agent/agent_0.py
@@ -543,7 +543,7 @@ def control_cb(self, topic, msg):
         self._log.debug_1('control msg %s: %s', topic, msg)
 
         cmd = msg['cmd']
-        arg = msg['arg']
+        arg = msg.get('arg')
 
         self._log.debug('pilot command: %s: %s', cmd, arg)
         self._prof.prof('cmd', msg="%s : %s" %  (cmd, arg), uid=self._pid)

diff --git a/src/radical/pilot/utils/component.py b/src/radical/pilot/utils/component.py
@@ -351,7 +351,7 @@ def _control_cb(self, topic, msg):
         self._log.debug_5('command incoming: %s', msg)
 
         cmd = msg['cmd']
-        arg = msg['arg']
+        arg = msg.get('arg')
 
         if cmd == 'cancel_tasks':
 
@@ -419,7 +419,7 @@ def _handle_rpc_msg(self, msg):
 
         if msg.cmd not in self._rpc_handlers:
             # this RPC message is *silently* ignored
-            self._log.debug('no rpc handler for [%s])', msg.cmd)
+            self._log.debug('no rpc handler for [%s]', msg.cmd)
             return
 
         rpc_handler, addr = self._rpc_handlers[msg.cmd]
@@ -428,7 +428,6 @@ def _handle_rpc_msg(self, msg):
             self._log.debug('ignore rpc handler for [%s] [%s])', msg, addr)
             return
 
-
         try:
             self._log.debug('rpc handler for %s: %s',
                             msg.cmd, self._rpc_handlers[msg.cmd])
@@ -452,10 +451,10 @@ def _handle_rpc_msg(self, msg):
             sys.stdout = bakout
             sys.stderr = bakerr
 
-        rpc_rep = RPCResultMessage(rpc_req=msg, val=val, out=out, err=err, exc=exc)
-        self._log.debug_3('rpc reply: %s', rpc_rep)
+        rpc_res = RPCResultMessage(rpc_req=msg, val=val, out=out, err=err, exc=exc)
+        self._log.debug_3('rpc reply: %s', rpc_res)
 
-        self.publish(rpc.CONTROL_PUBSUB, rpc_rep)
+        self.publish(rpc.CONTROL_PUBSUB, rpc_res)
 
 
     # --------------------------------------------------------------------------

diff --git a/tests/unit_tests/test_agent_0/test_agent_0.py b/tests/unit_tests/test_agent_0/test_agent_0.py
@@ -84,7 +84,7 @@ def _prepenv_effect(env_name, env_spec):
         agent_cmp.publish      = mock.MagicMock(side_effect=_publish_effect)
         agent_cmp._prepare_env = mock.MagicMock(side_effect=_prepenv_effect)
 
-        agent_cmp._rpc_handlers = {'prepare_env': agent_cmp._prepare_env}
+        agent_cmp._rpc_handlers = {'prepare_env': (agent_cmp._prepare_env, None)}
 
         msg = {'cmd': 'test',
                'arg': {'uid': 'rpc.0000',