def test_handle_unregistered_msg_type(self):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
msg = Message(source="worker", type="foo")
reply = mgr._handle_unregistered_msg_type(msg)
assert mgr.logger.error.call_count == 1
assert Message.deserialize(reply) == Message(
source="manager", type="unregistered message type")
def test_checklist_cleared_msg_type(self):
mgr = manager.NowcastManager()
mgr.checklist = {"foo": "bar"}
mgr.logger = Mock(name="logger")
mgr._write_checklist_to_disk = Mock(name="_write_checklist_to_disk")
reply = mgr._clear_checklist()
assert Message.deserialize(reply) == Message(source="manager",
type="checklist cleared")
def test_deserialize(self, source, msg_type, payload):
message = yaml.dump({
"source": source,
"type": msg_type,
"payload": payload
})
msg = Message.deserialize(message)
assert msg.source == source
assert msg.type == msg_type
assert msg.payload == payload
def test_missing_after_worker_function(self, m_importlib):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr._msg_registry = {"next workers module": "nowcast.next_workers"}
mgr._update_checklist = Mock(name="_update_checklist")
msg = Message(source="test_worker", type="success")
reply, next_workers = mgr._handle_continue_msg(msg)
assert Message.deserialize(reply) == Message(
source="manager", type="no after_worker function")
assert mgr.logger.critical.call_count == 1
def tell_manager(self, msg_type, payload=None):
"""Exchange messages with the nowcast manager process.
Message is composed of worker's name, msg_type, and payload.
Acknowledgement message from manager process is logged and returned.
:arg str msg_type: Key of the message type to send;
must be defined for worker name in the configuration
data structure.
:arg payload: Data object to send in the message;
e.g. dict containing worker's checklist of
accomplishments.
:returns: Acknowledgement message from manager process.
"""
try:
worker_msgs = self.config["message registry"]["workers"][self.name]
except (KeyError, TypeError):
raise WorkerError(
f"worker not found in {self.config.file} message registry: {self.name}"
)
try:
msg_words = worker_msgs[msg_type]
except (KeyError, TypeError):
raise WorkerError(
f"message type not found for {self.name} worker in {self.config.file} "
f"message registry: {msg_type}")
if self._parsed_args.debug:
self.logger.debug(
f"**debug mode** message that would have been sent to manager: ({msg_type} {msg_words})"
)
return
# Send message to nowcast manager
message = Message(self.name, msg_type, payload).serialize()
self._socket.send_string(message)
self.logger.debug(
f"sent message: ({msg_type}) {worker_msgs[msg_type]}",
extra={"logger_name": self.name},
)
# Wait for and process response
msg = self._socket.recv_string()
message = Message.deserialize(msg)
mgr_msgs = self.config["message registry"]["manager"]
try:
msg_words = mgr_msgs[message.type]
except KeyError:
raise WorkerError(
f"message type not found for manager in {self.config.file} message registry: {message.type}"
)
self.logger.debug(
f"received message from {message.source}: ({message.type}) {msg_words}",
extra={"logger_name": self.name},
)
return message
def test_reply(self, m_importlib):
mgr = manager.NowcastManager()
mgr._update_checklist = Mock(name="_update_checklist")
mgr._next_workers_module = Mock(
name="nowcast.next_workers",
after_test_worker=Mock(name="after_test_worker", return_value=[]),
)
msg = Message(source="test_worker", type="success")
reply, next_workers = mgr._handle_continue_msg(msg)
assert Message.deserialize(reply) == Message(source="manager",
type="ack")
def _message_handler(self, message):
"""Handle message from worker.
"""
msg = Message.deserialize(message)
if msg.source not in self._msg_registry["workers"]:
reply = self._handle_unregistered_worker_msg(msg)
return reply, []
if msg.type not in self._msg_registry["workers"][msg.source]:
reply = self._handle_unregistered_msg_type(msg)
return reply, []
self._log_received_msg(msg)
if msg.type == "clear checklist":
reply = self._clear_checklist()
return reply, []
if msg.type == "need":
reply = self._handle_need_msg(msg)
return reply, []
reply, next_workers = self._handle_continue_msg(msg)
return reply, next_workers
