Skip to content

Commit

Permalink
REF: Make attributes of MockMessage explicit
Browse files Browse the repository at this point in the history
  • Loading branch information
@cortadocodes
cortadocodes committed Jul 31, 2023
1 parent f752390 commit d24e6bd
Show file tree
Hide file tree
Showing 2 changed files with 59 additions and 39 deletions.
20 changes: 10 additions & 10 deletions octue/cloud/emulators/_pub_sub.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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()


Expand Down Expand Up @@ -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})>"
Expand Down Expand Up @@ -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
Expand Down
78 changes: 49 additions & 29 deletions tests/cloud/pub_sub/test_message_handler.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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(
Expand Down Expand Up @@ -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},
)
]

Expand Down

0 comments on commit d24e6bd

Please sign in to comment.