def test_tell_manager(self):
worker = NowcastWorker("test_worker", "description")
worker._parsed_args = Mock(debug=False)
worker._socket = Mock(name="_socket")
worker.logger = Mock(name="logger")
config = Config()
config.file = "nowcast.yaml"
worker.config._dict = {
"message registry": {
"manager": {
"ack": "message acknowledged"
},
"workers": {
"test_worker": {
"success": "successful test"
}
},
}
}
mgr_msg = Message(source="manager", type="ack")
worker._socket.recv_string.return_value = mgr_msg.serialize()
response = worker.tell_manager("success", "payload")
worker._socket.send_string.assert_called_once_with(
Message(source="test_worker", type="success",
payload="payload").serialize())
worker._socket.recv_string.assert_called_once_with()
assert worker.logger.debug.call_count == 2
assert response == mgr_msg
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")
class TestHandleNeedMsg:
"""Unit test for NowcastManager._handle_need_msg method.
"""
mgr = manager.NowcastManager()
mgr.checklist = {"info": "requested info"}
msg = Message(source="test_worker", type="need", payload="info")
reply = mgr._handle_need_msg(msg)
expected = Message("manager", "ack", payload="requested info").serialize()
assert reply == expected
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 test_unregistered_msg_type(self):
mgr = manager.NowcastManager()
mgr._msg_registry = {"workers": {"test_worker": {}}}
mgr._handle_unregistered_msg_type = Mock(
name="_handle_unregistered_msg_type")
mgr._log_received_msg = Mock(name="_log_received_msg")
msg = Message(source="test_worker", type="foo", payload=None)
reply, next_workers = mgr._message_handler(msg.serialize())
mgr._handle_unregistered_msg_type.assert_called_once_with(msg)
assert reply == mgr._handle_unregistered_msg_type()
assert next_workers == []
assert not mgr._log_received_msg.called
def test_need_msg(self):
mgr = manager.NowcastManager()
mgr._msg_registry = {"workers": {"test_worker": {"need": "Need info"}}}
mgr._handle_need_msg = Mock(
name="_handle_need_msg",
return_value=("ack msg w/ requested info in payload"),
)
mgr._log_received_msg = Mock(name="_log_received_msg")
msg = Message(source="test_worker", type="need", payload="info")
reply, next_workers = mgr._message_handler(msg.serialize())
assert mgr._log_received_msg.called
mgr._handle_need_msg.assert_called_once_with(msg)
assert reply == "ack msg w/ requested info in payload"
assert next_workers == []
def _handle_continue_msg(self, msg):
"""Handle success, failure, or crash message from worker by generating
list of subsequent workers to launch.
"""
if msg.payload is not None:
self._update_checklist(msg)
self._slack_notification(msg)
importlib.reload(self._next_workers_module)
worker = msg.source
try:
after_func = getattr(self._next_workers_module, f"after_{worker}")
except AttributeError:
self.logger.critical(
f"could not find after_{worker} in {self._msg_registry['next workers module']} module",
exc_info=True,
)
reply = Message(self.name, "no after_worker function").serialize()
return reply, []
next_workers = after_func(msg, self.config, self.checklist)
if len(next_workers) > 1 and isinstance(next_workers[-1], set):
next_workers, self._race_condition_mgmt[
"must finish"] = next_workers
self._race_condition_mgmt["then launch"] = []
self.logger.debug(
f"race condition management activated: {self._race_condition_mgmt}"
)
try:
self._race_condition_mgmt["must finish"].remove(worker)
self._race_condition_mgmt["then launch"].extend(next_workers)
next_workers.clear()
self.logger.debug(
f"{worker} finished and race condition management updated: {self._race_condition_mgmt}"
)
except (KeyError, ValueError):
# No race condition management in effect, or worker not in "must finish" list
pass
try:
if not self._race_condition_mgmt["must finish"]:
next_workers = self._race_condition_mgmt["then launch"]
self._race_condition_mgmt = {}
self.logger.debug(
f"race condition management ended; next workers released: {next_workers}"
)
except KeyError:
# No race condition management in effect
pass
reply = Message(self.name, "ack").serialize()
return reply, next_workers
def test_activate_race_condition_mgmt(self, m_importlib):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr._update_checklist = Mock(name="_update_checklist")
mgr._next_workers_module = Mock(
name="nowcast.next_workers",
after_test_worker=Mock(
name="after_download_weather",
return_value=(
[
NextWorker("get_NeahBay_ssh"),
NextWorker("grib_to_netcdf"),
NextWorker("download_live_ocean"),
],
{"grib_to_netcdf", "make_live_ocean_files"},
),
),
)
msg = Message(source="test_worker", type="success")
_, next_workers = mgr._handle_continue_msg(msg)
assert mgr._race_condition_mgmt == {
"must finish": {"grib_to_netcdf", "make_live_ocean_files"},
"then launch": [],
}
assert next_workers == [
NextWorker("get_NeahBay_ssh"),
NextWorker("grib_to_netcdf"),
NextWorker("download_live_ocean"),
]
def test_worker_checklist_keyerror(self):
mgr = manager.NowcastManager()
mgr._msg_registry = {"workers": {"test_worker": {}}}
msg = Message(source="test_worker",
type="success",
payload={"foo": "bar"})
with pytest.raises(KeyError):
mgr._update_checklist(msg)
def test_worker_not_in_notifications_list(self, m_post):
mgr = manager.NowcastManager()
mgr.config = {"slack notifications": {"SLACK_URL": []}}
msg = Message(source="test_worker", type="success")
with patch.dict(os.environ,
{"SLACK_URL": "https://hooks.slack.com/services/..."}):
mgr._slack_notification(msg)
assert not m_post.called
def test_notification_posted(self, m_post):
mgr = manager.NowcastManager()
mgr.config = {"slack notifications": {"SLACK_URL": ["test_worker"]}}
msg = Message(source="test_worker", type="success")
slack_url = "https://hooks.slack.com/services/..."
with patch.dict(os.environ, {"SLACK_URL": slack_url}):
mgr._slack_notification(msg)
m_post.assert_called_once_with(slack_url,
json={"text": "test_worker: success"})
def test_continue_msg(self):
mgr = manager.NowcastManager()
mgr._msg_registry = {
"workers": {
"test_worker": {
"success": "success"
}
}
}
mgr._handle_continue_msg = Mock(name="_handle_continue_msg",
return_value=("ack", "next_worker"))
mgr._log_received_msg = Mock(name="_log_received_msg")
msg = Message(source="test_worker", type="success", payload=None)
reply, next_workers = mgr._message_handler(msg.serialize())
assert mgr._log_received_msg.called
mgr._handle_continue_msg.assert_called_once_with(msg)
assert reply == "ack"
assert next_workers == "next_worker"
def test_no_slack_url_envvar(self, m_post):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr.config = {"slack notifications": {"SLACK_URL": []}}
msg = Message(source="test_worker", type="success")
mgr._slack_notification(msg)
mgr.logger.debug.assert_called_once_with(
"slack notification environment variable not found: SLACK_URL")
assert not m_post.called
def test_no_checklist_update_when_no_payload(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")
mgr._handle_continue_msg(msg)
assert not mgr._update_checklist.called
def _handle_unregistered_msg_type(self, msg):
"""Emit error message to log about a message type received from a worker
that is not included in the message registry.
"""
self.logger.error(
f"unregistered message type received from {msg.source} worker: {msg.type}",
extra={"worker_msg": msg},
)
reply = Message(self.name, "unregistered message type").serialize()
return reply
def test_slack_notification(self, m_importlib):
mgr = manager.NowcastManager()
mgr._slack_notification = Mock(name="_slack_notification")
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")
mgr._handle_continue_msg(msg)
assert mgr._slack_notification.called
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_reload_next_workers_module(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")
mgr._handle_continue_msg(msg)
m_importlib.reload.assert_called_once_with(mgr._next_workers_module)
def test_clear_checklist_msg(self):
mgr = manager.NowcastManager()
mgr._msg_registry = {
"workers": {
"test_worker": {
"clear checklist":
"request that manager clear system checklist"
}
}
}
mgr._clear_checklist = Mock(name="_clear_checklist",
return_value="checklist cleared")
mgr._log_received_msg = Mock(name="_log_received_msg")
msg = Message(source="test_worker",
type="clear checklist",
payload=None)
reply, next_workers = mgr._message_handler(msg.serialize())
assert mgr._log_received_msg.called
mgr._clear_checklist.assert_called_once_with()
assert reply == "checklist cleared"
assert next_workers == []
def test_one_next_worker_no_race_condition_mgmt(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=[NextWorker("another_test_worker")],
),
)
msg = Message(source="test_worker", type="success", payload=None)
reply, next_workers = mgr._handle_continue_msg(msg)
assert next_workers == [NextWorker("another_test_worker")]
def test_unregistered_manager_message_type(self):
worker = NowcastWorker("test_worker", "description")
worker._parsed_args = Mock(debug=False)
worker._socket = Mock(name="_socket")
worker.logger = Mock(name="logger")
config = Config()
config.file = "nowcast.yaml"
worker.config._dict = {
"message registry": {
"manager": {
"ack": "message acknowledged"
},
"workers": {
"test_worker": {
"success": "successful test"
}
},
}
}
mgr_msg = Message(source="manager", type="foo")
worker._socket.recv_string.return_value = mgr_msg.serialize()
with pytest.raises(WorkerError):
worker.tell_manager("success", "payload")
def test_yaml_dump_checklist_to_disk(self):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr.checklist = {"foo": "bar"}
mgr._write_checklist_to_disk = Mock(name="_write_checklist_to_disk")
mgr._msg_registry = {
"workers": {
"test_worker": {
"checklist key": "foo"
}
}
}
msg = Message(source="test_worker", type="success", payload="baz")
mgr._update_checklist(msg)
mgr._write_checklist_to_disk.assert_called_once_with()
def test_keyerror_adds_key_and_value(self):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr.checklist = {"foo": "bar"}
mgr._write_checklist_to_disk = Mock(name="_write_checklist_to_disk")
mgr._msg_registry = {
"workers": {
"test_worker": {
"checklist key": "fop"
}
}
}
msg = Message(source="test_worker", type="success", payload="baz")
mgr._update_checklist(msg)
assert mgr.checklist["fop"] == "baz"
def test_notification_posted_with_log_url(self, m_post):
mgr = manager.NowcastManager()
website_log_url = "https://salishsea.eos.ubc.ca/nemo/nowcast/logs/nowcast.log"
mgr.config = {
"slack notifications": {
"SLACK_URL": ["test_worker"],
"website log url": website_log_url,
}
}
msg = Message(source="test_worker", type="success")
slack_url = "https://hooks.slack.com/services/..."
with patch.dict(os.environ, {"SLACK_URL": slack_url}):
mgr._slack_notification(msg)
m_post.assert_called_once_with(
slack_url,
json={"text": f"test_worker: success\nLog: {website_log_url}"})
def test_log_received_msg(self):
mgr = manager.NowcastManager()
mgr.logger = Mock(name="logger")
mgr._msg_registry = {
"workers": {
"test_worker": {
"success": "worker succeeded"
}
}
}
msg = Message(source="test_worker", type="success")
mgr._log_received_msg(msg)
mgr.logger.debug.assert_called_once_with(
"received message from test_worker: (success) worker succeeded",
extra={"worker_msg": msg},
)
def test_log_info_msg(self):
mgr = manager.NowcastManager()
mgr.checklist = {"foo": "bar"}
mgr.logger = Mock(name="logger")
mgr._write_checklist_to_disk = Mock(name="_write_checklist_to_disk")
mgr._msg_registry = {
"workers": {
"test_worker": {
"checklist key": "foo"
}
}
}
msg = Message(source="test_worker", type="success", payload="baz")
mgr._update_checklist(msg)
mgr.logger.info.assert_called_once_with(
"checklist updated with [foo] items from test_worker worker",
extra={"worker_msg": msg},
)
def test_worker_not_in_race_condition_mgmt(self, m_importlib):
mgr = manager.NowcastManager()
mgr._update_checklist = Mock(name="_update_checklist")
mgr._next_workers_module = Mock(
name="nowcast.next_workers",
after_get_NeahBay_ssh=Mock(name="after_get_NeahBay_ssh",
return_value=[]),
)
mgr._race_condition_mgmt = {
"must finish": {"grib_to_netcdf", "make_live_ocean_files"},
"then launch": [],
}
msg = Message(source="get_NeahBay_ssh", type="success")
mgr._handle_continue_msg(msg)
assert mgr._race_condition_mgmt == {
"must finish": {"grib_to_netcdf", "make_live_ocean_files"},
"then launch": [],
}