def to_cps_subscribe_message(source: SequencedMessage) -> PubsubMessage:
source_pb = source._pb
out_pb = _to_cps_publish_message_proto(source_pb.message)
out_pb.publish_time.CopyFrom(source_pb.publish_time)
out = PubsubMessage()
out._pb = out_pb
return out
def callback(message: PubsubMessage):
message_data = message.data.decode("utf-8")
metadata = MessageMetadata.decode(message.message_id)
print(
f"Received {message_data} of ordering key {message.ordering_key} with id {metadata}."
)
message.ack()
def test_messages_received(subscriber: SubscriberImpl, async_subscriber,
message_callback, close_callback):
message1 = Message(PubsubMessage(message_id="1")._pb, "", 0, None)
message2 = Message(PubsubMessage(message_id="2")._pb, "", 0, None)
counter = Box[int]()
counter.val = 0
async def on_read() -> Message:
counter.val += 1
if counter.val == 1:
return message1
if counter.val == 2:
return message2
await sleep_forever()
async_subscriber.read.side_effect = on_read
results = Queue()
message_callback.side_effect = lambda m: results.put(m.message_id)
subscriber.add_close_callback(close_callback)
subscriber.__enter__()
assert results.get() == "1"
assert results.get() == "2"
subscriber.close()
async def test_iterator(
default_subscriber,
subscriber_factory,
multiplexed_client: AsyncSubscriberClientInterface,
):
read_queues = wire_queues(default_subscriber.read)
subscription = SubscriptionPath(1, CloudZone.parse("us-central1-a"), "abc")
message = Message(PubsubMessage(message_id="1")._pb, "", 0, None)
async with multiplexed_client:
iterator = await multiplexed_client.subscribe(
subscription, DISABLED_FLOW_CONTROL
)
subscriber_factory.assert_has_calls(
[call(subscription, None, DISABLED_FLOW_CONTROL)]
)
read_fut_1 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_1.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call()])
await read_queues.results.put(message)
assert await read_fut_1 is message
read_fut_2 = asyncio.ensure_future(iterator.__anext__())
assert not read_fut_2.done()
await read_queues.called.get()
default_subscriber.read.assert_has_calls([call(), call()])
await read_queues.results.put(FailedPrecondition(""))
with pytest.raises(FailedPrecondition):
await read_fut_2
default_subscriber.__aexit__.assert_called_once()
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_publish_invalid_event_time():
with pytest.raises(InvalidArgument):
from_cps_publish_message(
PubsubMessage(
attributes={PUBSUB_LITE_EVENT_TIME: "probably not an encoded proto"}
)
)
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
async def publish(self,
data: bytes,
ordering_key: str = "",
**attrs: Mapping[str, str]) -> str:
cps_message = PubsubMessage(data=data,
ordering_key=ordering_key,
attributes=attrs)
psl_message = from_cps_publish_message(cps_message)
return (await self._publisher.publish(psl_message)).encode()
async def test_delivery_from_multiple(subscriber, assigner,
subscriber_factory):
assign_queues = wire_queues(assigner.get_assignment)
async with subscriber:
await assign_queues.called.get()
sub1 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub2 = mock_async_context_manager(
MagicMock(spec=AsyncSingleSubscriber))
sub1_queues = wire_queues(sub1.read)
sub2_queues = wire_queues(sub2.read)
subscriber_factory.side_effect = (
lambda partition: sub1 if partition == Partition(1) else sub2)
await assign_queues.results.put({Partition(1), Partition(2)})
await sub1_queues.results.put(
Message(PubsubMessage(message_id="1")._pb, "", 0, None))
await sub2_queues.results.put(
Message(PubsubMessage(message_id="2")._pb, "", 0, None))
message_ids: Set[str] = set()
message_ids.add((await subscriber.read()).message_id)
message_ids.add((await subscriber.read()).message_id)
assert message_ids == {"1", "2"}
def publish_message(publisher, topic_name, message):
project_id = os.environ[GOOGLE_CLOUD_PROJECT]
if not topic_name:
raise Exception(
f"env var TOPIC_DISPATCHER is badly set. It's current value is: {topic_name}"
)
if not project_id:
raise Exception(
f"env var GOOGLE_CLOUD_PROJECT is badly set. It's current value is: {project_id}"
)
topic_path = f"projects/{project_id}/topics/{topic_name}"
message_json = json.dumps(message)
encoded_message = message_json.encode("utf-8")
message_obj = PubsubMessage(data=encoded_message)
future = publisher.publish(topic=topic_path, messages=[message_obj])
try:
return future.result()
except AttributeError as ae:
print(f"AttributeError on future.result. future is {future}")
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
def to_cps_publish_message(source: PubSubMessage) -> PubsubMessage:
out = PubsubMessage()
try:
out.ordering_key = source.key.decode("utf-8")
except UnicodeError:
raise InvalidArgument(
"Received an unparseable message with a non-utf8 key.")
if PUBSUB_LITE_EVENT_TIME in source.attributes:
raise InvalidArgument(
"Special timestamp attribute exists in wire message. Unable to parse message."
)
out.data = source.data
for key, values in source.attributes.items():
out.attributes[key] = _parse_attributes(values)
if "event_time" in source:
out.attributes[PUBSUB_LITE_EVENT_TIME] = encode_attribute_event_time(
source.event_time)
return out
def create_fake_message(
message_data,
ack_id="ACK_ID",
message_id="MESSAGE_ID",
timestamp_epoch_seconds=int(time.time())
) -> ReceivedMessage:
publish_time_timestamp = Timestamp()
publish_time_timestamp.seconds = timestamp_epoch_seconds
pubsub_message = PubsubMessage()
pubsub_message.data = message_data.encode("UTF-8")
pubsub_message.message_id = message_id
pubsub_message.publish_time = publish_time_timestamp
pubsub_message.ordering_key = ""
received_message = ReceivedMessage()
received_message.ack_id = ack_id
received_message.message = pubsub_message
received_message.delivery_attempt = 0
return received_message
def transformer():
result = MagicMock(spec=MessageTransformer)
result.transform.side_effect = lambda source: PubsubMessage(message_id=str(
source.cursor.offset))
return result
def to_cps_publish_message(source: PubSubMessage) -> PubsubMessage:
out = PubsubMessage()
out._pb = _to_cps_publish_message_proto(source._pb)
return out