def test_msg_manager_maybe_extend(
mocker, monkeypatch, msg_manager, last_extended, deadline_extended, caplog,
):
mock_time = mocker.Mock()
monkeypatch.setattr(pmm, "time", mock_time)
mock_time.monotonic.return_value = 15
mock_ext_deadline = mocker.Mock()
monkeypatch.setattr(msg_manager, "extend_deadline", mock_ext_deadline)
kmsg = pmm.PubSubKlioMessage(ack_id=1, kmsg_id=2)
kmsg.last_extended = last_extended
kmsg.ext_duration = 10 # threshold = 0.8 * 10 = 8s
msg_manager._maybe_extend(kmsg)
mgr_logs = [
c
for c in caplog.records
if c.name == msg_manager.mgr_logger.name and c.levelno == logging.DEBUG
]
if deadline_extended:
mock_ext_deadline.assert_called_once_with(kmsg)
assert 0 == len(mgr_logs)
else:
mock_ext_deadline.assert_not_called()
assert 1 == len(mgr_logs)
def test_msg_manager_remove_raises(msg_manager, monkeypatch, caplog):
ack_id, kmsg_id = 1, "2"
psk_msg1 = pmm.PubSubKlioMessage(ack_id, kmsg_id)
msg_manager._client.acknowledge.side_effect = Exception("oh no")
msg_manager.remove(psk_msg1)
assert 2 == len(caplog.records)
def test_msg_manager_remove(mocker, monkeypatch, msg_manager, caplog):
ack_id, kmsg_id = 1, "2"
psk_msg1 = pmm.PubSubKlioMessage(ack_id, kmsg_id)
msg_manager.remove(psk_msg1)
msg_manager._client.acknowledge.assert_called_once_with(
msg_manager._sub_name, [ack_id]
)
assert 1 == len(caplog.records)
def test_msg_manager_add(mocker, monkeypatch, msg_manager, caplog):
extend_deadline = mocker.Mock()
monkeypatch.setattr(msg_manager, "extend_deadline", extend_deadline)
mock_event = mocker.Mock()
monkeypatch.setattr(pmm.threading, "Event", mock_event)
pmsg1 = _get_pubsub_message("2")
pmsg2 = _get_pubsub_message("4")
psk_msg1 = pmm.PubSubKlioMessage(ack_id=1, kmsg_id="2")
psk_msg2 = pmm.PubSubKlioMessage(ack_id=3, kmsg_id="4")
msg_manager.add(ack_id=1, raw_pubsub_message=pmsg1)
msg_manager.add(ack_id=3, raw_pubsub_message=pmsg2)
assert 4 == msg_manager.executor.submit.call_count
assert 2 == len(caplog.records)
assert _compare_objects_dicts(
psk_msg1, pmm.ENTITY_ID_TO_ACK_ID[psk_msg1.kmsg_id]
)
assert _compare_objects_dicts(
psk_msg2, pmm.ENTITY_ID_TO_ACK_ID[psk_msg2.kmsg_id]
)
def test_klio_ack_input_msg(mocker, monkeypatch):
ack_id, kmsg_id = 1, "2"
psk_msg1 = pmm.PubSubKlioMessage(ack_id, kmsg_id)
monkeypatch.setitem(pmm.ENTITY_ID_TO_ACK_ID, psk_msg1.kmsg_id, psk_msg1)
with beam_test_pipeline.TestPipeline() as p:
(
p
| beam.Create([kmsg_id])
| beam.Map(_generate_kmsg)
| beam.ParDo(helpers.KlioAckInputMessage())
| beam.Map(_assert_expected_msg)
)
assert not pmm.ENTITY_ID_TO_ACK_ID.get("2")
def test_convert_raw_pubsub_message(mocker, monkeypatch, msg_manager):
mock_event = mocker.Mock()
monkeypatch.setattr(pmm.threading, "Event", mock_event)
exp_message = pmm.PubSubKlioMessage("ack_id1", "kmsg_id1")
kmsg = klio_pb2.KlioMessage()
kmsg.data.element = b"kmsg_id1"
kmsg_bytes = kmsg.SerializeToString()
pmsg = beam_pubsub.PubsubMessage(data=kmsg_bytes, attributes={})
act_message = msg_manager._convert_raw_pubsub_message("ack_id1", pmsg)
# comparing class attributes (via __dict__) since we'd need to implement
# __eq__ on the PubSubKlioMessage class, but doing so would make it un-
# hashable. Which can be addressed, but this just seems easier for now.
assert _compare_objects_dicts(exp_message, act_message)
def test_msg_manager_extend_deadline(mocker, msg_manager, duration, caplog):
kmsg = pmm.PubSubKlioMessage(ack_id=1, kmsg_id=2)
kmsg.extend = mocker.Mock()
msg_manager.extend_deadline(kmsg, duration)
exp_duration = (
duration if duration else msg_manager.DEFAULT_DEADLINE_EXTENSION
)
exp_req = {
"subscription": msg_manager._sub_name,
"ack_ids": [kmsg.ack_id],
"ack_deadline_seconds": exp_duration,
}
msg_manager._client.modify_ack_deadline.assert_called_once_with(**exp_req)
assert 1 == len(caplog.records)
kmsg.extend.assert_called_once_with(exp_duration)
def pubsub_klio_msg():
m = pmm.PubSubKlioMessage(ack_id=1, kmsg_id=2)
return m
def test_msg_manager_extend_deadline_raises(msg_manager, caplog):
kmsg = pmm.PubSubKlioMessage(ack_id=1, kmsg_id=2)
msg_manager._client.modify_ack_deadline.side_effect = Exception("oh no")
msg_manager.extend_deadline(kmsg, 12)
assert 2 == len(caplog.records)