def test_offsets(self):
event = self.event
message: Message[KafkaPayload] = Message(
Partition(Topic("events"), 456),
123,
KafkaPayload(
None, json.dumps((0, "insert", event)).encode("utf-8")
), # event doesn't really matter
datetime.now(),
)
test_worker = ConsumerWorker(
self.dataset,
producer=FakeConfluentKafkaProducer(),
replacements_topic=Topic(
enforce_table_writer(self.dataset)
.get_stream_loader()
.get_replacement_topic_spec()
.topic_name
),
metrics=self.metrics,
)
batch = [test_worker.process_message(message)]
test_worker.flush_batch(batch)
assert self.clickhouse.execute(
"SELECT project_id, event_id, offset, partition FROM %s" % self.table
) == [(self.event["project_id"], self.event["event_id"], 123, 456)]
def test_synchronized_consumer_worker_crash_before_assignment(
broker: Broker[KafkaPayload], ) -> None:
topic = Topic("topic")
commit_log_topic = Topic("commit-log")
broker.create_topic(topic, partitions=1)
broker.create_topic(commit_log_topic, partitions=1)
poll_called = Event()
class BrokenConsumerException(Exception):
pass
class BrokenConsumer(LocalConsumer[KafkaPayload]):
def poll(
self,
timeout: Optional[float] = None
) -> Optional[Message[KafkaPayload]]:
try:
raise BrokenConsumerException()
finally:
poll_called.set()
consumer = broker.get_consumer("consumer")
commit_log_consumer: Consumer[KafkaPayload] = BrokenConsumer(
broker, "commit-log-consumer")
with pytest.raises(BrokenConsumerException):
Consumer[KafkaPayload] = SynchronizedConsumer(
consumer,
commit_log_consumer,
commit_log_topic=commit_log_topic,
commit_log_groups={"leader"},
)
def test_skip_too_old(self):
test_worker = ConsumerWorker(
self.dataset,
producer=FakeConfluentKafkaProducer(),
replacements_topic=Topic(
enforce_table_writer(self.dataset)
.get_stream_loader()
.get_replacement_topic_spec()
.topic_name
),
metrics=self.metrics,
)
event = self.event
old_timestamp = datetime.utcnow() - timedelta(days=300)
old_timestamp_str = old_timestamp.strftime("%Y-%m-%dT%H:%M:%S.%fZ")
event["datetime"] = old_timestamp_str
event["data"]["datetime"] = old_timestamp_str
event["data"]["received"] = int(calendar.timegm(old_timestamp.timetuple()))
message: Message[KafkaPayload] = Message(
Partition(Topic("events"), 1),
42,
KafkaPayload(None, json.dumps((0, "insert", event)).encode("utf-8")),
datetime.now(),
)
assert test_worker.process_message(message) is None
def test_filter() -> None:
next_step = Mock()
def test_function(message: Message[bool]) -> bool:
return message.payload
filter_step = FilterStep(test_function, next_step)
fail_message = Message(Partition(Topic("topic"), 0), 0, False,
datetime.now())
with assert_does_not_change(lambda: next_step.submit.call_count, 0):
filter_step.submit(fail_message)
pass_message = Message(Partition(Topic("topic"), 0), 0, True,
datetime.now())
with assert_changes(lambda: next_step.submit.call_count, 0, 1):
filter_step.submit(pass_message)
assert next_step.submit.call_args == call(pass_message)
with assert_changes(lambda: next_step.poll.call_count, 0, 1):
filter_step.poll()
with assert_changes(lambda: next_step.close.call_count, 0,
1), assert_changes(lambda: next_step.join.call_count,
0, 1):
filter_step.join()
def test_synchronized_consumer_worker_crash_after_assignment(
broker: Broker[KafkaPayload], ) -> None:
topic = Topic("topic")
commit_log_topic = Topic("commit-log")
broker.create_topic(topic, partitions=1)
broker.create_topic(commit_log_topic, partitions=1)
poll_called = Event()
class BrokenConsumerException(Exception):
pass
class BrokenConsumer(LocalConsumer[KafkaPayload]):
def poll(
self,
timeout: Optional[float] = None
) -> Optional[Message[KafkaPayload]]:
if not self.tell():
return super().poll(timeout)
else:
try:
raise BrokenConsumerException()
finally:
poll_called.set()
consumer: Consumer[KafkaPayload] = broker.get_consumer("consumer")
commit_log_consumer: Consumer[KafkaPayload] = BrokenConsumer(
broker, "commit-log-consumer")
synchronized_consumer: Consumer[KafkaPayload] = SynchronizedConsumer(
consumer,
commit_log_consumer,
commit_log_topic=commit_log_topic,
commit_log_groups={"leader"},
)
assert poll_called.wait(1.0) is True
# If the worker thread has exited without a close request, calling ``poll``
# should raise an error that originated from the worker thread.
with pytest.raises(RuntimeError) as e:
synchronized_consumer.poll(0.0)
assert type(e.value.__cause__) is BrokenConsumerException
# If a close request has been sent, the normal runtime error due to the
# closed consumer should be raised instead.
synchronized_consumer.close()
with pytest.raises(RuntimeError) as e:
synchronized_consumer.poll(0.0)
assert type(e.value.__cause__) is not BrokenConsumerException
def test_synchronized_consumer_handles_end_of_partition(
broker: Broker[KafkaPayload], ) -> None:
topic = Topic("topic")
commit_log_topic = Topic("commit-log")
broker.create_topic(topic, partitions=1)
broker.create_topic(commit_log_topic, partitions=1)
consumer = broker.get_consumer("consumer", enable_end_of_partition=True)
producer = broker.get_producer()
commit_log_consumer = broker.get_consumer("commit-log-consumer")
messages = [
producer.produce(topic, KafkaPayload(None, f"{i}".encode("utf8"),
[])).result(1.0) for i in range(2)
]
synchronized_consumer: Consumer[KafkaPayload] = SynchronizedConsumer(
consumer,
commit_log_consumer,
commit_log_topic=commit_log_topic,
commit_log_groups={"leader"},
)
with closing(synchronized_consumer):
synchronized_consumer.subscribe([topic])
wait_for_consumer(
commit_log_consumer,
producer.produce(
commit_log_topic,
commit_codec.encode(
Commit("leader", Partition(topic, 0),
messages[0].next_offset), ),
).result(),
)
assert synchronized_consumer.poll(0) == messages[0]
# If the commit log consumer does not handle EOF, it will have crashed
# here and will never return the next message.
wait_for_consumer(
commit_log_consumer,
producer.produce(
commit_log_topic,
commit_codec.encode(
Commit("leader", Partition(topic, 0),
messages[1].next_offset), ),
).result(),
)
assert synchronized_consumer.poll(0) == messages[1]
def test_synchronized_consumer_handles_end_of_partition() -> None:
topic = Topic("topic")
commit_log_topic = Topic("commit-log")
broker: DummyBroker[int] = DummyBroker()
broker.create_topic(topic, partitions=1)
consumer: Consumer[int] = DummyConsumer(broker, "consumer")
producer: Producer[int] = DummyProducer(broker)
messages = [producer.produce(topic, i).result(1.0) for i in range(2)]
commit_log_broker: DummyBroker[Commit] = DummyBroker()
commit_log_broker.create_topic(commit_log_topic, partitions=1)
commit_log_consumer: Consumer[Commit] = DummyConsumer(
commit_log_broker, "commit-log-consumer", enable_end_of_partition=True)
commit_log_producer: Producer[Commit] = DummyProducer(commit_log_broker)
synchronized_consumer: Consumer[int] = SynchronizedConsumer(
consumer,
commit_log_consumer,
commit_log_topic=commit_log_topic,
commit_log_groups={"leader"},
)
with closing(synchronized_consumer):
synchronized_consumer.subscribe([topic])
wait_for_consumer(
commit_log_consumer,
commit_log_producer.produce(
commit_log_topic,
Commit("leader", Partition(topic, 0),
messages[0].get_next_offset()),
).result(),
)
assert synchronized_consumer.poll(0) == messages[0]
# If the commit log consumer does not handle EOF, it will have crashed
# here and will never return the next message.
wait_for_consumer(
commit_log_consumer,
commit_log_producer.produce(
commit_log_topic,
Commit("leader", Partition(topic, 0),
messages[1].get_next_offset()),
).result(),
)
assert synchronized_consumer.poll(0) == messages[1]
def test_encoding_producer(broker: Broker[str]) -> None:
topic = Topic("test")
broker.create_topic(topic, 1)
class ReverseEncoder(Encoder[str, str]):
def encode(self, value: str) -> str:
return "".join(value[::-1])
producer = ProducerEncodingWrapper(broker.get_producer(), ReverseEncoder())
decoded_message = producer.produce(topic, "hello").result()
assert decoded_message.payload == "hello"
consumer = broker.get_consumer("group")
consumer.subscribe([topic])
encoded_message = consumer.poll()
assert encoded_message is not None
# The payload returned by the consumer should not be decoded.
assert encoded_message.payload == "olleh"
# All other attributes should be the same.
for attribute in set(Message.__slots__) - {"payload"}:
assert getattr(encoded_message,
attribute) == getattr(decoded_message, attribute)
def eventstream(*, dataset: Dataset):
ensure_table_exists(dataset)
record = json.loads(http_request.data)
version = record[0]
if version != 2:
raise RuntimeError("Unsupported protocol version: %s" % record)
message: Message[KafkaPayload] = Message(
Partition(Topic("topic"), 0),
0,
KafkaPayload(None, http_request.data),
datetime.now(),
)
type_ = record[1]
metrics = DummyMetricsBackend()
if type_ == "insert":
from snuba.consumer import ConsumerWorker
worker = ConsumerWorker(dataset, metrics=metrics)
else:
from snuba.replacer import ReplacerWorker
worker = ReplacerWorker(clickhouse_rw, dataset, metrics=metrics)
processed = worker.process_message(message)
if processed is not None:
batch = [processed]
worker.flush_batch(batch)
return ("ok", 200, {"Content-Type": "text/plain"})
def test_stream_processor_termination_on_error() -> None:
topic = Topic("test")
consumer = mock.Mock()
consumer.poll.return_value = Message(Partition(topic, 0), 0, 0,
datetime.now())
exception = NotImplementedError("error")
strategy = mock.Mock()
strategy.submit.side_effect = exception
factory = mock.Mock()
factory.create.return_value = strategy
processor: StreamProcessor[int] = StreamProcessor(consumer, topic, factory,
TestingMetricsBackend())
assignment_callback = consumer.subscribe.call_args.kwargs["on_assign"]
assignment_callback({Partition(topic, 0): 0})
with pytest.raises(Exception) as e, assert_changes(
lambda: strategy.terminate.call_count, 0,
1), assert_changes(lambda: consumer.close.call_count, 0, 1):
processor.run()
assert e.value == exception
def test_delete_groups_insert(self):
self.event["project_id"] = self.project_id
self.event["group_id"] = 1
self.write_raw_events(self.event)
assert self._issue_count(self.project_id) == [{"count": 1, "group_id": 1}]
timestamp = datetime.now(tz=pytz.utc)
project_id = self.project_id
message: Message[KafkaPayload] = Message(
Partition(Topic("replacements"), 1),
42,
KafkaPayload(
None,
json.dumps(
(
2,
"end_delete_groups",
{
"project_id": project_id,
"group_ids": [1],
"datetime": timestamp.strftime(PAYLOAD_DATETIME_FORMAT),
},
)
).encode("utf-8"),
),
datetime.now(),
)
processed = self.replacer.process_message(message)
self.replacer.flush_batch([processed])
assert self._issue_count(self.project_id) == []
def test_transform() -> None:
next_step = Mock()
def transform_function(message: Message[int]) -> int:
return message.payload * 2
transform_step = TransformStep(transform_function, next_step)
original_message = Message(Partition(Topic("topic"), 0), 0, 1,
datetime.now())
with assert_changes(lambda: next_step.submit.call_count, 0, 1):
transform_step.submit(original_message)
assert next_step.submit.call_args == call(
Message(
original_message.partition,
original_message.offset,
transform_function(original_message),
original_message.timestamp,
))
with assert_changes(lambda: next_step.poll.call_count, 0, 1):
transform_step.poll()
with assert_changes(lambda: next_step.close.call_count, 0,
1), assert_changes(lambda: next_step.join.call_count,
0, 1):
transform_step.join()
def test_batch_size(self, broker: Broker[int]) -> None:
topic = Topic("topic")
broker.create_topic(topic, partitions=1)
producer = broker.get_producer()
for i in [1, 2, 3]:
producer.produce(topic, i).result()
consumer = broker.get_consumer("group")
worker = FakeWorker()
batching_consumer = StreamProcessor(
consumer,
topic,
BatchProcessingStrategyFactory(
worker=worker,
max_batch_size=2,
max_batch_time=100,
metrics=DummyMetricsBackend(strict=True),
),
)
for _ in range(3):
batching_consumer._run_once()
batching_consumer._shutdown()
assert worker.processed == [1, 2, 3]
assert worker.flushed == [[1, 2]]
assert consumer.commit_offsets_calls == 1
assert consumer.close_calls == 1
def test_send_message(
self,
value: str,
expected: Optional[ProcessedMessage],
) -> None:
storage = get_storage("groupedmessages")
snapshot_id = uuid1()
transact_data = TransactionData(xmin=100,
xmax=200,
xip_list=[120, 130])
worker = SnapshotAwareWorker(
storage=storage,
producer=FakeConfluentKafkaProducer(),
snapshot_id=str(snapshot_id),
transaction_data=transact_data,
replacements_topic=None,
metrics=DummyMetricsBackend(strict=True),
)
message: Message[KafkaPayload] = Message(
Partition(Topic("topic"), 0),
1,
KafkaPayload(
None,
value.encode("utf-8"),
[("table", "sentry_groupedmessage".encode())],
),
datetime.now(),
)
ret = worker.process_message(message)
assert ret == expected
def _wrap(self, msg: str) -> Message[KafkaPayload]:
return Message(
Partition(Topic("replacements"), 0),
0,
KafkaPayload(None, json.dumps(msg).encode("utf-8")),
datetime.now(),
)
def test_produce_replacement_messages(self):
producer = FakeConfluentKafkaProducer()
test_worker = ConsumerWorker(
self.dataset,
producer=producer,
replacements_topic=Topic(
enforce_table_writer(self.dataset)
.get_stream_loader()
.get_replacement_topic_spec()
.topic_name
),
metrics=self.metrics,
)
test_worker.flush_batch(
[
ProcessedMessage(
action=ProcessorAction.REPLACE, data=[("1", {"project_id": 1})],
),
ProcessedMessage(
action=ProcessorAction.REPLACE, data=[("2", {"project_id": 2})],
),
]
)
assert [(m._topic, m._key, m._value) for m in producer.messages] == [
("event-replacements", b"1", b'{"project_id": 1}'),
("event-replacements", b"2", b'{"project_id": 2}'),
]
def __delivery_callback(
self,
future: Future[Message[KafkaPayload]],
payload: KafkaPayload,
error: KafkaError,
message: ConfluentMessage,
) -> None:
if error is not None:
future.set_exception(TransportError(error))
else:
try:
timestamp_type, timestamp_value = message.timestamp()
if timestamp_type is TIMESTAMP_NOT_AVAILABLE:
raise ValueError("timestamp not available")
future.set_result(
Message(
Partition(Topic(message.topic()), message.partition()),
message.offset(),
payload,
datetime.utcfromtimestamp(timestamp_value / 1000.0),
)
)
except Exception as error:
future.set_exception(error)
def __make_msg(self, partition: int, offset: int, payload: str,
headers: Headers) -> Message[KafkaPayload]:
return Message(
partition=Partition(Topic("topic"), partition),
offset=offset,
payload=KafkaPayload(b"key", payload.encode(), headers),
timestamp=datetime(2019, 6, 19, 6, 46, 28),
)
def test_batch_time(self, mock_time: Any, broker: Broker[int]) -> None:
topic = Topic("topic")
broker.create_topic(topic, partitions=1)
producer = broker.get_producer()
consumer = broker.get_consumer("group")
worker = FakeWorker()
metrics = DummyMetricsBackend(strict=True)
batching_consumer = StreamProcessor(
consumer,
topic,
BatchProcessingStrategyFactory(
worker=worker,
max_batch_size=100,
max_batch_time=2000,
metrics=metrics,
),
metrics=metrics,
)
mock_time.return_value = time.mktime(
datetime(2018, 1, 1, 0, 0, 0).timetuple())
for i in [1, 2, 3]:
producer.produce(topic, i).result()
for _ in range(3):
batching_consumer._run_once()
mock_time.return_value = time.mktime(
datetime(2018, 1, 1, 0, 0, 1).timetuple())
for i in [4, 5, 6]:
producer.produce(topic, i).result()
for _ in range(3):
batching_consumer._run_once()
mock_time.return_value = time.mktime(
datetime(2018, 1, 1, 0, 0, 5).timetuple())
for i in [7, 8, 9]:
producer.produce(topic, i).result()
for _ in range(3):
batching_consumer._run_once()
batching_consumer._shutdown()
assert worker.processed == [1, 2, 3, 4, 5, 6, 7, 8, 9]
assert worker.flushed == [[1, 2, 3, 4, 5, 6]]
assert consumer.commit_offsets_calls == 1
assert consumer.close_calls == 1
def test_unaligned_offset(self) -> None:
topic = Topic(uuid.uuid1().hex)
partition = Partition(topic, 0)
self.storage.create_topic(topic, 1)
message = self.storage.produce(partition, 1, datetime.now())
invalid_offset = message.offset + 4
assert message.next_offset > invalid_offset > message.offset
with pytest.raises(InvalidChecksum):
self.storage.consume(partition, invalid_offset)
def test_storage(self) -> None:
topic = Topic(uuid.uuid1().hex)
partitions = 3
self.storage.create_topic(topic, partitions)
assert [*self.storage.list_topics()] == [topic]
assert self.storage.get_partition_count(topic) == partitions
with pytest.raises(TopicExists):
self.storage.create_topic(topic, partitions)
with pytest.raises(TopicDoesNotExist):
self.storage.get_partition_count(Topic("invalid"))
with pytest.raises(TopicDoesNotExist):
self.storage.consume(Partition(Topic("invalid"), 0), 0)
with pytest.raises(TopicDoesNotExist):
self.storage.produce(Partition(Topic("invalid"), 0), 0,
datetime.now())
with pytest.raises(PartitionDoesNotExist):
self.storage.consume(Partition(topic, -1), 0)
with pytest.raises(PartitionDoesNotExist):
self.storage.consume(Partition(topic, partitions + 1), 0)
with pytest.raises(PartitionDoesNotExist):
self.storage.produce(Partition(topic, -1), 0, datetime.now())
with pytest.raises(PartitionDoesNotExist):
self.storage.produce(Partition(topic, partitions + 1), 0,
datetime.now())
self.storage.delete_topic(topic)
with pytest.raises(TopicDoesNotExist):
self.storage.delete_topic(topic)
def test_parallel_transform_worker_apply() -> None:
messages = [
Message(
Partition(Topic("test"), 0),
i,
KafkaPayload(None, b"\x00" * size, None),
datetime.now(),
) for i, size in enumerate([1000, 1000, 2000, 4000])
]
with SharedMemoryManager() as smm:
input_block = smm.SharedMemory(8192)
assert input_block.size == 8192
input_batch = MessageBatch(input_block)
for message in messages:
input_batch.append(message)
assert len(input_batch) == 4
output_block = smm.SharedMemory(4096)
assert output_block.size == 4096
index, output_batch = parallel_transform_worker_apply(
transform_payload_expand,
input_batch,
output_block,
)
# The first batch should be able to fit 2 messages.
assert index == 2
assert len(output_batch) == 2
index, output_batch = parallel_transform_worker_apply(
transform_payload_expand,
input_batch,
output_block,
index,
)
# The second batch should be able to fit one message.
assert index == 3
assert len(output_batch) == 1
# The last message is too large to fit in the batch.
with pytest.raises(ValueTooLarge):
parallel_transform_worker_apply(
transform_payload_expand,
input_batch,
output_block,
index,
)
def decode(self, value: KafkaPayload) -> Commit:
key = value.key
if not isinstance(key, bytes):
raise TypeError("payload key must be a bytes object")
val = value.value
if not isinstance(val, bytes):
raise TypeError("payload value must be a bytes object")
topic_name, partition_index, group = key.decode("utf-8").split(":", 3)
offset = int(val.decode("utf-8"))
return Commit(group, Partition(Topic(topic_name),
int(partition_index)), offset)
def test_delete_tag_promoted_insert(self):
self.event["project_id"] = self.project_id
self.event["group_id"] = 1
self.event["data"]["tags"].append(["browser.name", "foo"])
self.event["data"]["tags"].append(["notbrowser", "foo"])
self.write_raw_events(self.event)
project_id = self.project_id
def _issue_count(total=False):
return json.loads(
self.app.post(
"/query",
data=json.dumps({
"project": [project_id],
"aggregations": [["count()", "", "count"]],
"conditions": [["tags[browser.name]", "=", "foo"]]
if not total else [],
"groupby": ["group_id"],
}),
).data)["data"]
assert _issue_count() == [{"count": 1, "group_id": 1}]
assert _issue_count(total=True) == [{"count": 1, "group_id": 1}]
timestamp = datetime.now(tz=pytz.utc)
message: Message[KafkaPayload] = Message(
Partition(Topic("replacements"), 1),
42,
KafkaPayload(
None,
json.dumps((
2,
"end_delete_tag",
{
"project_id": project_id,
"tag": "browser.name",
"datetime":
timestamp.strftime(PAYLOAD_DATETIME_FORMAT),
},
)).encode("utf-8"),
),
datetime.now(),
)
processed = self.replacer.process_message(message)
self.replacer.flush_batch([processed])
assert _issue_count() == []
assert _issue_count(total=True) == [{"count": 1, "group_id": 1}]
def get_topic(self, partitions: int = 1) -> Iterator[Topic]:
name = f"test-{uuid.uuid1().hex}"
client = AdminClient(self.configuration)
[[key, future]] = client.create_topics(
[NewTopic(name, num_partitions=partitions,
replication_factor=1)]).items()
assert key == name
assert future.result() is None
try:
yield Topic(name)
finally:
[[key, future]] = client.delete_topics([name]).items()
assert key == name
assert future.result() is None
def get_messages(events_file) -> Sequence[Message[KafkaPayload]]:
"Create a fake Kafka message for each JSON event in the file."
messages: MutableSequence[Message[KafkaPayload]] = []
raw_events = open(events_file).readlines()
for raw_event in raw_events:
messages.append(
Message(
Partition(Topic("events"), 1),
0,
KafkaPayload(None, raw_event.encode("utf-8")),
datetime.now(),
),
)
return messages
def test_commit_log_consumer(self) -> None:
# XXX: This would be better as an integration test (or at least a test
# against an abstract Producer interface) instead of against a test against
# a mock.
commit_log_producer = FakeConfluentKafkaProducer()
consumer: KafkaConsumer[int] = KafkaConsumerWithCommitLog(
{
**self.configuration,
"auto.offset.reset": "earliest",
"enable.auto.commit": "false",
"enable.auto.offset.store": "false",
"enable.partition.eof": "true",
"group.id": "test",
"session.timeout.ms": 10000,
},
codec=self.codec,
producer=commit_log_producer,
commit_log_topic=Topic("commit-log"),
)
with self.get_topic() as topic, closing(consumer) as consumer:
consumer.subscribe([topic])
with closing(self.get_producer()) as producer:
producer.produce(topic, 0).result(5.0)
message = consumer.poll(
10.0) # XXX: getting the subscription is slow
assert isinstance(message, Message)
consumer.stage_offsets(
{message.partition: message.get_next_offset()})
assert consumer.commit_offsets() == {
Partition(topic, 0): message.get_next_offset()
}
assert len(commit_log_producer.messages) == 1
commit_message = commit_log_producer.messages[0]
assert commit_message.topic() == "commit-log"
assert CommitCodec().decode(
KafkaPayload(commit_message.key(),
commit_message.value())) == Commit(
"test", Partition(topic, 0),
message.get_next_offset())
def assignment_callback(
consumer: ConfluentConsumer, partitions: Sequence[ConfluentTopicPartition]
) -> None:
self.__state = KafkaConsumerState.ASSIGNING
try:
assignment: MutableSequence[ConfluentTopicPartition] = []
for partition in self.__consumer.committed(partitions):
if partition.offset >= 0:
assignment.append(partition)
elif partition.offset == OFFSET_INVALID:
assignment.append(
self.__resolve_partition_starting_offset(partition)
)
else:
raise ValueError("received unexpected offset")
offsets: MutableMapping[Partition, int] = {
Partition(Topic(i.topic), i.partition): i.offset for i in assignment
}
self.__seek(offsets)
# Ensure that all partitions are resumed on assignment to avoid
# carrying over state from a previous assignment.
self.__consumer.resume(
[
ConfluentTopicPartition(
partition.topic.name, partition.index, offset
)
for partition, offset in offsets.items()
]
)
for partition in offsets:
self.__paused.discard(partition)
except Exception:
self.__state = KafkaConsumerState.ERROR
raise
try:
if on_assign is not None:
on_assign(offsets)
finally:
self.__state = KafkaConsumerState.CONSUMING
def test_collect() -> None:
step_factory = Mock()
step_factory.return_value = inner_step = Mock()
commit_function = Mock()
partition = Partition(Topic("topic"), 0)
messages = message_generator(partition, 0)
collect_step = CollectStep(step_factory, commit_function, 2, 60)
# A batch should be started the first time the step receives a message.
with assert_changes(lambda: step_factory.call_count, 0, 1):
collect_step.poll()
collect_step.submit(next(messages)) # offset 0
# Subsequent messages should reuse the existing batch, ...
with assert_does_not_change(lambda: step_factory.call_count, 1):
collect_step.poll()
collect_step.submit(next(messages)) # offset 1
# ...until we hit the batch size limit.
with assert_changes(lambda: inner_step.close.call_count,
0, 1), assert_changes(
lambda: inner_step.join.call_count, 0,
1), assert_changes(
lambda: commit_function.call_count, 0, 1):
collect_step.poll()
assert commit_function.call_args == call({partition: 2})
step_factory.return_value = inner_step = Mock()
# The next message should create a new batch.
with assert_changes(lambda: step_factory.call_count, 1, 2):
collect_step.submit(next(messages))
with assert_changes(lambda: inner_step.close.call_count, 0, 1):
collect_step.close()
with assert_changes(lambda: inner_step.join.call_count, 0,
1), assert_changes(lambda: commit_function.call_count,
1, 2):
collect_step.join()
def test_message_batch() -> None:
partition = Partition(Topic("test"), 0)
with SharedMemoryManager() as smm:
block = smm.SharedMemory(4096)
assert block.size == 4096
message = Message(partition, 0, KafkaPayload(None, b"\x00" * 4000,
None), datetime.now())
batch: MessageBatch[KafkaPayload] = MessageBatch(block)
with assert_changes(lambda: len(batch), 0, 1):
batch.append(message)
assert batch[0] == message
assert list(batch) == [message]
with assert_does_not_change(lambda: len(batch),
1), pytest.raises(ValueTooLarge):
batch.append(message)