From d24e6bde6ecc7a43f7ee9c21ee237d88dd1603b8 Mon Sep 17 00:00:00 2001 From: cortadocodes Date: Mon, 31 Jul 2023 17:23:01 +0100 Subject: [PATCH] REF: Make attributes of `MockMessage` explicit --- octue/cloud/emulators/_pub_sub.py | 20 +++--- tests/cloud/pub_sub/test_message_handler.py | 78 +++++++++++++-------- 2 files changed, 59 insertions(+), 39 deletions(-) diff --git a/octue/cloud/emulators/_pub_sub.py b/octue/cloud/emulators/_pub_sub.py index be29caf5d..da617e5e6 100644 --- a/octue/cloud/emulators/_pub_sub.py +++ b/octue/cloud/emulators/_pub_sub.py @@ -146,7 +146,7 @@ def publish(self, topic, data, retry=None, **attributes): :return MockFuture: """ subscription_name = ".".join((get_pub_sub_resource_name(topic), ANSWERS_NAMESPACE, attributes["question_uuid"])) - SUBSCRIPTIONS[subscription_name].append(MockMessage(data=data, **attributes)) + SUBSCRIPTIONS[subscription_name].append(MockMessage(data=data, attributes=attributes)) return MockFuture() @@ -249,11 +249,9 @@ class MockMessage: :return None: """ - def __init__(self, data, **attributes): + def __init__(self, data, attributes=None): self.data = data - self.attributes = {} - for key, value in attributes.items(): - self.attributes[key] = value + self.attributes = attributes or {} def __repr__(self): return f"<{type(self).__name__}(data={self.data!r})>" @@ -345,11 +343,13 @@ def ask( {"input_values": input_values, "input_manifest": input_manifest, "children": children}, cls=OctueJSONEncoder, ).encode(), - question_uuid=question_uuid, - forward_logs=subscribe_to_logs, - octue_sdk_version=parent_sdk_version, - allow_save_diagnostics_data_on_crash=allow_save_diagnostics_data_on_crash, - is_question=True, + attributes={ + "question_uuid": question_uuid, + "forward_logs": subscribe_to_logs, + "octue_sdk_version": parent_sdk_version, + "allow_save_diagnostics_data_on_crash": allow_save_diagnostics_data_on_crash, + "is_question": True, + }, ) ) except Exception as e: # noqa diff --git a/tests/cloud/pub_sub/test_message_handler.py b/tests/cloud/pub_sub/test_message_handler.py index 9d2860428..958b9f3b7 100644 --- a/tests/cloud/pub_sub/test_message_handler.py +++ b/tests/cloud/pub_sub/test_message_handler.py @@ -70,10 +70,10 @@ def test_in_order_messages_are_handled_in_order(self): ) messages = [ - MockMessage(json.dumps({"type": "test"}).encode(), message_number=0), - MockMessage(json.dumps({"type": "test"}).encode(), message_number=1), - MockMessage(json.dumps({"type": "test"}).encode(), message_number=2), - MockMessage(json.dumps({"type": "finish-test"}).encode(), message_number=3), + MockMessage(json.dumps({"type": "test"}).encode(), attributes={"message_number": 0}), + MockMessage(json.dumps({"type": "test"}).encode(), attributes={"message_number": 1}), + MockMessage(json.dumps({"type": "test"}).encode(), attributes={"message_number": 2}), + MockMessage(json.dumps({"type": "finish-test"}).encode(), attributes={"message_number": 3}), ] with patch( @@ -95,10 +95,13 @@ def test_out_of_order_messages_are_handled_in_order(self): ) messages = [ - MockMessage(data=json.dumps({"type": "test", "order": 1}).encode(), message_number=1), - MockMessage(data=json.dumps({"type": "test", "order": 2}).encode(), message_number=2), - MockMessage(data=json.dumps({"type": "test", "order": 0}).encode(), message_number=0), - MockMessage(data=json.dumps({"type": "finish-test", "order": 3}).encode(), message_number=3), + MockMessage(data=json.dumps({"type": "test", "order": 1}).encode(), attributes={"message_number": 1}), + MockMessage(data=json.dumps({"type": "test", "order": 2}).encode(), attributes={"message_number": 2}), + MockMessage(data=json.dumps({"type": "test", "order": 0}).encode(), attributes={"message_number": 0}), + MockMessage( + data=json.dumps({"type": "finish-test", "order": 3}).encode(), + attributes={"message_number": 3}, + ), ] with patch( @@ -133,10 +136,22 @@ def test_out_of_order_messages_with_end_message_first_are_handled_in_order(self) "octue.cloud.pub_sub.service.OrderedMessageHandler._pull_and_enqueue_message", new=MockMessagePuller( messages=[ - MockMessage(data=json.dumps({"type": "finish-test", "order": 3}).encode(), message_number=3), - MockMessage(data=json.dumps({"type": "test", "order": 1}).encode(), message_number=1), - MockMessage(data=json.dumps({"type": "test", "order": 2}).encode(), message_number=2), - MockMessage(data=json.dumps({"type": "test", "order": 0}).encode(), message_number=0), + MockMessage( + data=json.dumps({"type": "finish-test", "order": 3}).encode(), + attributes={"message_number": 3}, + ), + MockMessage( + data=json.dumps({"type": "test", "order": 1}).encode(), + attributes={"message_number": 1}, + ), + MockMessage( + data=json.dumps({"type": "test", "order": 2}).encode(), + attributes={"message_number": 2}, + ), + MockMessage( + data=json.dumps({"type": "test", "order": 0}).encode(), + attributes={"message_number": 0}, + ), ], message_handler=message_handler, ).pull, @@ -165,9 +180,12 @@ def test_no_timeout(self): ) messages = [ - MockMessage(data=json.dumps({"type": "test", "order": 0}).encode(), message_number=0), - MockMessage(data=json.dumps({"type": "test", "order": 1}).encode(), message_number=1), - MockMessage(data=json.dumps({"type": "finish-test", "order": 2}).encode(), message_number=2), + MockMessage(data=json.dumps({"type": "test", "order": 0}).encode(), attributes={"message_number": 0}), + MockMessage(data=json.dumps({"type": "test", "order": 1}).encode(), attributes={"message_number": 1}), + MockMessage( + data=json.dumps({"type": "finish-test", "order": 2}).encode(), + attributes={"message_number": 2}, + ), ] with patch( @@ -201,11 +219,11 @@ def test_delivery_acknowledgement(self): "delivery_time": "2021-11-17 17:33:59.717428", } ).encode(), - message_number=0, + attributes={"message_number": 0}, ), MockMessage( data=json.dumps({"type": "result", "output_values": None, "output_manifest": None}).encode(), - message_number=1, + attributes={"message_number": 1}, ), ], message_handler=message_handler, @@ -266,11 +284,11 @@ def test_error_not_raised_if_heartbeat_has_been_received_in_maximum_allowed_inte "delivery_time": "2021-11-17 17:33:59.717428", }, ).encode(), - message_number=0, + attributes={"message_number": 0}, ), MockMessage( data=json.dumps({"type": "result", "output_values": None, "output_manifest": None}).encode(), - message_number=1, + attributes={"message_number": 1}, ), ], message_handler=message_handler, @@ -318,10 +336,13 @@ def test_handler_can_skip_first_n_messages_if_missed(self): message_handler._earliest_message_number_received = 2 messages = [ - MockMessage(data=json.dumps({"type": "test", "order": 2}).encode(), message_number=2), - MockMessage(data=json.dumps({"type": "test", "order": 3}).encode(), message_number=3), - MockMessage(data=json.dumps({"type": "test", "order": 4}).encode(), message_number=4), - MockMessage(data=json.dumps({"type": "finish-test", "order": 5}).encode(), message_number=5), + MockMessage(data=json.dumps({"type": "test", "order": 2}).encode(), attributes={"message_number": 2}), + MockMessage(data=json.dumps({"type": "test", "order": 3}).encode(), attributes={"message_number": 3}), + MockMessage(data=json.dumps({"type": "test", "order": 4}).encode(), attributes={"message_number": 4}), + MockMessage( + data=json.dumps({"type": "finish-test", "order": 5}).encode(), + attributes={"message_number": 5}, + ), ] with patch( @@ -351,10 +372,10 @@ def test_later_missing_messages_cannot_be_skipped(self): ) messages = [ - MockMessage(json.dumps({"type": "test", "order": 0}).encode(), message_number=0), - MockMessage(json.dumps({"type": "test", "order": 1}).encode(), message_number=1), - MockMessage(json.dumps({"type": "test", "order": 2}).encode(), message_number=2), - MockMessage(json.dumps({"type": "finish-test", "order": 5}).encode(), message_number=5), + MockMessage(json.dumps({"type": "test", "order": 0}).encode(), attributes={"message_number": 0}), + MockMessage(json.dumps({"type": "test", "order": 1}).encode(), attributes={"message_number": 1}), + MockMessage(json.dumps({"type": "test", "order": 2}).encode(), attributes={"message_number": 2}), + MockMessage(json.dumps({"type": "finish-test", "order": 5}).encode(), attributes={"message_number": 5}), ] with patch( @@ -402,8 +423,7 @@ def test_pull_and_enqueue_message(self): SUBSCRIPTIONS[mock_subscription.name] = [ MockMessage( data=json.dumps(mock_message).encode(), - is_question=False, - message_number=0, + attributes={"is_question": False, "message_number": 0}, ) ]