async def test_simple_publish(mock_publishers, mock_policies, mock_watcher, publisher):
mock_watcher.get_partition_count.return_value = 2
async with publisher:
mock_policies[2].route.return_value = Partition(1)
mock_publishers[Partition(1)].publish.return_value = "a"
await publisher.publish(PubSubMessage())
mock_policies[2].route.assert_called_with(PubSubMessage())
mock_publishers[Partition(1)].publish.assert_called()
def test_subscribe_transform_correct():
expected = PubsubMessage(
data=b"xyz",
ordering_key="def",
attributes={
"x": "abc",
"y": "abc",
PUBSUB_LITE_EVENT_TIME: encode_attribute_event_time(
Timestamp(seconds=55).ToDatetime()
),
},
publish_time=Timestamp(seconds=10),
)
result = to_cps_subscribe_message(
SequencedMessage(
message=PubSubMessage(
data=b"xyz",
key=b"def",
event_time=Timestamp(seconds=55),
attributes={
"x": AttributeValues(values=[b"abc"]),
"y": AttributeValues(values=[b"abc"]),
},
),
publish_time=Timestamp(seconds=10),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
assert result == expected
def test_wrapped_successful():
wrapped = add_id_to_cps_subscribe_transformer(
Partition(1), MessageTransformer.of_callable(to_cps_subscribe_message)
)
expected = PubsubMessage(
data=b"xyz",
ordering_key="def",
attributes={
"x": "abc",
"y": "abc",
PUBSUB_LITE_EVENT_TIME: encode_attribute_event_time(
Timestamp(seconds=55).ToDatetime()
),
},
message_id=MessageMetadata(Partition(1), Cursor(offset=10)).encode(),
publish_time=Timestamp(seconds=10),
)
result = wrapped.transform(
SequencedMessage(
message=PubSubMessage(
data=b"xyz",
key=b"def",
event_time=Timestamp(seconds=55),
attributes={
"x": AttributeValues(values=[b"abc"]),
"y": AttributeValues(values=[b"abc"]),
},
),
publish_time=Timestamp(seconds=10),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
assert result == expected
def test_invalid_subscribe_transform_key():
with pytest.raises(InvalidArgument):
to_cps_subscribe_message(
SequencedMessage(
message=PubSubMessage(key=NOT_UTF8),
publish_time=Timestamp(),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
def test_invalid_subscribe_contains_non_utf8_attributes():
with pytest.raises(InvalidArgument):
to_cps_subscribe_message(
SequencedMessage(
message=PubSubMessage(
key=b"def", attributes={"xyz": AttributeValues(values=[NOT_UTF8])}
),
publish_time=Timestamp(seconds=10),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
def from_cps_publish_message(source: PubsubMessage) -> PubSubMessage:
source_pb = source._pb
out = PubSubMessage()
out_pb = out._pb
if PUBSUB_LITE_EVENT_TIME in source_pb.attributes:
out_pb.event_time.CopyFrom(
_decode_attribute_event_time_proto(
source_pb.attributes[PUBSUB_LITE_EVENT_TIME]))
out_pb.data = source_pb.data
out_pb.key = source_pb.ordering_key.encode("utf-8")
for key, value in source_pb.attributes.items():
if key != PUBSUB_LITE_EVENT_TIME:
out_pb.attributes[key].values.append(value.encode("utf-8"))
return out
async def test_publish_after_increase(
mock_publishers, mock_policies, mock_watcher, publisher
):
get_queues = wire_queues(mock_watcher.get_partition_count)
await get_queues.results.put(2)
async with publisher:
get_queues.called.get_nowait()
mock_policies[2].route.return_value = Partition(1)
mock_publishers[Partition(1)].publish.return_value = "a"
await publisher.publish(PubSubMessage())
mock_policies[2].route.assert_called_with(PubSubMessage())
mock_publishers[Partition(1)].publish.assert_called()
await get_queues.called.get()
await get_queues.results.put(3)
await get_queues.called.get()
mock_policies[3].route.return_value = Partition(2)
mock_publishers[Partition(2)].publish.return_value = "a"
await publisher.publish(PubSubMessage())
mock_policies[3].route.assert_called_with(PubSubMessage())
mock_publishers[Partition(2)].publish.assert_called()
def test_routing_cases():
policy = DefaultRoutingPolicy(num_partitions=29)
json_list = []
with open(os.path.join(os.path.dirname(__file__),
"routing_tests.json")) as f:
for line in f:
if not line.startswith("//"):
json_list.append(line)
loaded = json.loads("\n".join(json_list))
target = {bytes(k, "utf-8"): Partition(v) for k, v in loaded.items()}
result = {}
for key in target:
result[key] = policy.route(PubSubMessage(key=key))
assert result == target
def test_invalid_subscribe_contains_magic_attribute():
with pytest.raises(InvalidArgument):
to_cps_subscribe_message(
SequencedMessage(
message=PubSubMessage(
key=b"def",
attributes={
PUBSUB_LITE_EVENT_TIME: AttributeValues(values=[b"abc"])
},
),
publish_time=Timestamp(seconds=10),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
def from_cps_publish_message(source: PubsubMessage) -> PubSubMessage:
out = PubSubMessage()
if PUBSUB_LITE_EVENT_TIME in source.attributes:
out.event_time = decode_attribute_event_time(
source.attributes[PUBSUB_LITE_EVENT_TIME])
out.data = source.data
out.key = source.ordering_key.encode("utf-8")
for key, value in source.attributes.items():
if key != PUBSUB_LITE_EVENT_TIME:
out.attributes[key] = AttributeValues(
values=[value.encode("utf-8")])
return out
def test_wrapped_sets_id_error():
wrapped = add_id_to_cps_subscribe_transformer(
Partition(1),
MessageTransformer.of_callable(lambda x: PubsubMessage(message_id="a")),
)
with pytest.raises(InvalidArgument):
wrapped.transform(
SequencedMessage(
message=PubSubMessage(
data=b"xyz",
key=b"def",
event_time=Timestamp(seconds=55),
attributes={
"x": AttributeValues(values=[b"abc"]),
"y": AttributeValues(values=[b"abc"]),
},
),
publish_time=Timestamp(seconds=10),
cursor=Cursor(offset=10),
size_bytes=10,
)
)
def test_publish_valid_transform():
now = datetime.datetime.now()
expected = PubSubMessage(
data=b"xyz",
key=b"def",
event_time=now,
attributes={
"x": AttributeValues(values=[b"abc"]),
"y": AttributeValues(values=[b"abc"]),
},
)
result = from_cps_publish_message(
PubsubMessage(
data=b"xyz",
ordering_key="def",
attributes={
"x": "abc",
"y": "abc",
PUBSUB_LITE_EVENT_TIME: encode_attribute_event_time(now),
},
)
)
assert result == expected