async def test_ack(subscriber: AsyncSingleSubscriber, underlying, transformer,
ack_set_tracker):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message_1 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
message_2 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=10)
underlying.read.return_value = message_1
read_1: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
assert read_1.message_id == "1"
underlying.read.return_value = message_2
read_2: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1), call(2)])
assert read_2.message_id == "2"
read_2.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(2)])
read_1.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(2), call(1)])
def test_add_remove():
batcher = FlowControlBatcher()
batcher.add(FlowControlRequest(allowed_bytes=10, allowed_messages=3))
restart_1 = batcher.request_for_restart()
assert restart_1.allowed_bytes == 10
assert restart_1.allowed_messages == 3
batcher.on_messages(
[SequencedMessage(size_bytes=2),
SequencedMessage(size_bytes=3)])
restart_2 = batcher.request_for_restart()
assert restart_2.allowed_bytes == 5
assert restart_2.allowed_messages == 1
async def test_nack_calls_ack(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
nack_handler,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
def on_nack(nacked: PubsubMessage, ack: Callable[[], None]):
assert nacked.message_id == "1"
ack()
nack_handler.on_nack.side_effect = on_nack
read.nack()
await ack_called_queue.get()
await ack_result_queue.put(None)
ack_set_tracker.ack.assert_has_calls([call(1)])
async def test_ack_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue)
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
read.ack()
await ack_called_queue.get()
ack_set_tracker.ack.assert_has_calls([call(1)])
await ack_result_queue.put(FailedPrecondition("Bad ack"))
async def sleep_forever():
await asyncio.sleep(float("inf"))
underlying.read.side_effect = sleep_forever
with pytest.raises(FailedPrecondition):
await subscriber.read()
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 _wrap_message(self, message: SequencedMessage.meta.pb) -> Message:
# Rewrap in the proto-plus-python wrapper for passing to the transform
rewrapped = SequencedMessage()
rewrapped._pb = message
cps_message = self._transformer.transform(rewrapped)
offset = message.cursor.offset
ack_id_str = _AckId(self._ack_generation_id, offset).encode()
self._ack_set_tracker.track(offset)
self._messages_by_ack_id[ack_id_str] = _SizedMessage(
cps_message, message.size_bytes)
wrapped_message = Message(
cps_message._pb,
ack_id=ack_id_str,
delivery_attempt=0,
request_queue=self._queue,
)
return wrapped_message
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,
)
)
async def test_handle_reset(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
ack_called_queue = asyncio.Queue()
ack_result_queue = asyncio.Queue()
ack_set_tracker.ack.side_effect = make_queue_waiter(
ack_called_queue, ack_result_queue
)
async with subscriber:
message_1 = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message_1
read_1: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
assert read_1.message_id == "1"
assert read_1.ack_id == ack_id(0, 1)
await subscriber.handle_reset()
ack_set_tracker.clear_and_commit.assert_called_once()
# Message ACKed after reset. Its flow control tokens are refilled
# but offset not committed (verified below after message 2).
read_1.ack()
message_2 = SequencedMessage(cursor=Cursor(offset=2), size_bytes=10)
underlying.read.return_value = message_2
read_2: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1), call(2)])
assert read_2.message_id == "2"
assert read_2.ack_id == ack_id(1, 2)
read_2.ack()
await ack_called_queue.get()
await ack_result_queue.put(None)
underlying.allow_flow.assert_has_calls(
[
call(FlowControlRequest(allowed_messages=1000, allowed_bytes=1000,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=5,)),
call(FlowControlRequest(allowed_messages=1, allowed_bytes=10,)),
]
)
ack_set_tracker.ack.assert_has_calls([call(2)])
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,
)
)
async def test_track_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
):
async with subscriber:
ack_set_tracker.track.side_effect = FailedPrecondition("Bad track")
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
with pytest.raises(FailedPrecondition):
await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
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,
)
)
async def test_nack_failure(
subscriber: SinglePartitionSingleSubscriber,
underlying,
transformer,
ack_set_tracker,
nack_handler,
):
async with subscriber:
message = SequencedMessage(cursor=Cursor(offset=1), size_bytes=5)
underlying.read.return_value = message
read: Message = await subscriber.read()
ack_set_tracker.track.assert_has_calls([call(1)])
nack_handler.on_nack.side_effect = FailedPrecondition("Bad nack")
read.nack()
async def sleep_forever():
await asyncio.sleep(float("inf"))
underlying.read.side_effect = sleep_forever
with pytest.raises(FailedPrecondition):
await subscriber.read()
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,
)
)
async def test_failed_transform(subscriber, underlying, transformer):
async with subscriber:
transformer.transform.side_effect = FailedPrecondition("Bad message")
underlying.read.return_value = SequencedMessage()
with pytest.raises(FailedPrecondition):
await subscriber.read()