def test_topic_parse_correctly_records():
headers = [
("header_1", "välûe%_1ù"),
("header_2", "välûe%_°2ù"),
("header_3", "välûe%_$*3ù"),
]
records_by_partition = {
TopicPartition(topic="topic", partition="partition"): [
ConsumerRecord(
topic="topic",
partition="partition",
offset=0,
timestamp=1562566,
timestamp_type=0,
key=None,
value=kirby_value_serializer("value"),
headers=[(header[0], kirby_value_serializer(header[1]))
for header in headers],
checksum=None,
serialized_key_size=None,
serialized_value_size=None,
serialized_header_size=None,
)
]
}
parsed_records = parse_records(records_by_partition, raw_records=True)
assert parsed_records[0].headers == {
header[0]: header[1]
for header in headers
}
def print_offsets(cls, partition_to_offset, tp_timestamps):
topics = {}
for tp, offset_timestamp in six.iteritems(partition_to_offset):
if tp.topic not in topics:
topics[tp.topic] = {}
topics[tp.topic][tp.partition] = offset_timestamp
topics = OrderedDict(sorted(topics.items(), key=lambda k: k[0]))
for topic in six.iterkeys(topics):
topics[topic] = OrderedDict(
sorted(topics[topic].items(), key=lambda k: k[0]))
print("Topic Name: {}".format(topic))
for partition, offset_timestamp in six.iteritems(topics[topic]):
print("\tPartition ID: {}".format(partition), )
offset = "not found"
timestamp = tp_timestamps[TopicPartition(topic, partition)]
if offset_timestamp is not None:
offset = offset_timestamp.offset
timestamp = offset_timestamp.timestamp
date = datetime.fromtimestamp(
timestamp / 1000.0,
tz=pytz.timezone("US/Pacific"),
).strftime("%Y-%m-%d %H:%M:%S %Z")
print(
"\t\tTimestamp: {timestamp} ({date})".format(
timestamp=timestamp,
date=date,
), )
print("\t\tOffset: {offset}".format(offset=offset))
def test_async_producer_not_leader(self):
for i in range(10):
self.queue.put((TopicPartition("test", i), "msg %i", "key %i"))
# Mock offsets counter for closure
offsets = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))
self.client.is_first_time = True
def send_side_effect(reqs, *args, **kwargs):
if self.client.is_first_time:
self.client.is_first_time = False
return [ProduceResponsePayload(req.topic, req.partition,
NotLeaderForPartitionError.errno, -1)
for req in reqs]
responses = []
for req in reqs:
offset = offsets[req.topic][req.partition]
offsets[req.topic][req.partition] += len(req.messages)
responses.append(
ProduceResponsePayload(req.topic, req.partition, 0, offset)
)
return responses
self.client.send_produce_request.side_effect = send_side_effect
self._run_process(2)
# the queue should be void at the end of the test
self.assertEqual(self.queue.empty(), True)
# there should be 5 non-void calls: 1st failed batch of 3 msgs
# + 3 batches of 3 msgs each + 1 batch of 1 msg = 1 + 3 + 1 = 5
self.assertEqual(self.client.send_produce_request.call_count, 5)
def test_first_send_failed(self):
# lets create a queue and add 10 messages for 10 different partitions
# to show how retries should work ideally
for i in range(10):
self.queue.put((TopicPartition("test", i), "msg %i", "key %i"))
# Mock offsets counter for closure
offsets = collections.defaultdict(lambda: collections.defaultdict(lambda: 0))
self.client.is_first_time = True
def send_side_effect(reqs, *args, **kwargs):
if self.client.is_first_time:
self.client.is_first_time = False
return [FailedPayloadsError(req) for req in reqs]
responses = []
for req in reqs:
offset = offsets[req.topic][req.partition]
offsets[req.topic][req.partition] += len(req.messages)
responses.append(
ProduceResponsePayload(req.topic, req.partition, 0, offset)
)
return responses
self.client.send_produce_request.side_effect = send_side_effect
self._run_process(2)
# the queue should be void at the end of the test
self.assertEqual(self.queue.empty(), True)
# there should be 5 non-void calls: 1st failed batch of 3 msgs
# plus 3 batches of 3 msgs each + 1 batch of 1 message
self.assertEqual(self.client.send_produce_request.call_count, 5)
def _find_coordinator_callback(self, consumer_group, response):
"""Callback that takes a FindCoordinatorResponse and issues an OffsetFetchRequest for the group.
consumer_group must be manually passed in because it is not present in the response, but we need it in order to
associate these offsets to the proper consumer group.
The OffsetFetchRequest is scoped to the topics and partitions that are specified in the check config. If
topics are unspecified, it will fetch all known offsets for that consumer group. Similiarly, if the partitions
are unspecified for a topic listed in the config, offsets are fetched for all the partitions within that topic.
"""
coordinator_id = self._kafka_client._find_coordinator_id_process_response(
response)
topics = self._consumer_groups[consumer_group]
if topics is None:
topic_partitions = None # None signals to fetch all known offsets for the consumer group
else:
# transform [("t1", [1, 2])] into [TopicPartition("t1", 1), TopicPartition("t1", 2)]
topic_partitions = []
for topic, partitions in topics.items():
if partitions is None: # If partitions aren't specified, fetch all partitions in the topic
partitions = self._kafka_client._client.cluster.partitions_for_topic(
topic)
topic_partitions.extend(
[TopicPartition(topic, p) for p in partitions])
single_group_offsets_future = self._kafka_client._list_consumer_group_offsets_send_request(
group_id=consumer_group,
group_coordinator_id=coordinator_id,
partitions=topic_partitions)
single_group_offsets_future.add_callback(
self._single_group_offsets_callback, consumer_group)
self._consumer_futures.append(single_group_offsets_future)
def _handle_offset_fetch_response(self, future, response):
offsets = {}
for topic, partitions in response.topics:
for partition, offset, metadata, error_code in partitions:
tp = TopicPartition(topic, partition)
error_type = Errors.for_code(error_code)
if error_type is not Errors.NoError:
error = error_type()
log.debug("Group %s failed to fetch offset for partition"
" %s: %s", self.group_id, tp, error)
if error_type is Errors.GroupLoadInProgressError:
# just retry
future.failure(error)
elif error_type is Errors.NotCoordinatorForGroupError:
# re-discover the coordinator and retry
self.coordinator_dead(error_type())
future.failure(error)
elif error_type is Errors.UnknownTopicOrPartitionError:
log.warning("OffsetFetchRequest -- unknown topic %s"
" (have you committed any offsets yet?)",
topic)
continue
else:
log.error("Unknown error fetching offsets for %s: %s",
tp, error)
future.failure(error)
return
elif offset >= 0:
# record the position with the offset
# (-1 indicates no committed offset to fetch)
offsets[tp] = OffsetAndMetadata(offset, metadata)
else:
log.debug("Group %s has no committed offset for partition"
" %s", self.group_id, tp)
future.success(offsets)
def test_kafka_consumer__blocking(self):
TIMEOUT_MS = 500
consumer = self.kafka_consumer(auto_offset_reset='earliest',
consumer_timeout_ms=TIMEOUT_MS)
# Manual assignment avoids overhead of consumer group mgmt
consumer.unsubscribe()
consumer.assign([TopicPartition(self.topic, 0)])
# Ask for 5 messages, nothing in queue, block 500ms
with Timer() as t:
with self.assertRaises(StopIteration):
msg = next(consumer)
self.assertGreaterEqual(t.interval, TIMEOUT_MS / 1000.0 )
self.send_messages(0, range(0, 10))
# Ask for 5 messages, 10 in queue. Get 5 back, no blocking
messages = set()
with Timer() as t:
for i in range(5):
msg = next(consumer)
messages.add((msg.partition, msg.offset))
self.assertEqual(len(messages), 5)
self.assertLess(t.interval, TIMEOUT_MS / 1000.0 )
# Ask for 10 messages, get 5 back, block 500ms
messages = set()
with Timer() as t:
with self.assertRaises(StopIteration):
for i in range(10):
msg = next(consumer)
messages.add((msg.partition, msg.offset))
self.assertEqual(len(messages), 5)
self.assertGreaterEqual(t.interval, TIMEOUT_MS / 1000.0 )
def _handle_metadata_update(self, cluster):
# if we encounter any unauthorized topics, raise an exception
if cluster.unauthorized_topics:
raise Errors.TopicAuthorizationFailedError(
cluster.unauthorized_topics)
if self._subscription.subscribed_pattern:
topics = []
for topic in cluster.topics(
self.config['exclude_internal_topics']):
if self._subscription.subscribed_pattern.match(topic):
topics.append(topic)
if set(topics) != self._subscription.subscription:
self._subscription.change_subscription(topics)
self._client.set_topics(
self._subscription.group_subscription())
# check if there are any changes to the metadata which should trigger
# a rebalance
if self._subscription.partitions_auto_assigned():
metadata_snapshot = self._build_metadata_snapshot(
self._subscription, cluster)
if self._metadata_snapshot != metadata_snapshot:
self._metadata_snapshot = metadata_snapshot
# If we haven't got group coordinator support,
# just assign all partitions locally
if self._auto_assign_all_partitions():
self._subscription.assign_from_subscribed([
TopicPartition(topic, partition)
for topic in self._subscription.subscription
for partition in self._metadata_snapshot[topic]
])
def test_conflicting_previous_assignments(mocker):
cluster = create_cluster(mocker, topics={'t'}, topics_partitions={0, 1})
subscriptions = {
'C1': {'t'},
'C2': {'t'},
}
member_metadata = {}
for member, topics in six.iteritems(subscriptions):
# assume both C1 and C2 have partition 1 assigned to them in generation 1
member_metadata[member] = build_metadata(
topics, [TopicPartition('t', 0),
TopicPartition('t', 0)], 1)
assignment = StickyPartitionAssignor.assign(cluster, member_metadata)
verify_validity_and_balance(subscriptions, assignment)
def test_refresh_committed_offsets_if_needed(mocker, coordinator):
mocker.patch.object(ConsumerCoordinator,
'fetch_committed_offsets',
return_value={
TopicPartition('foobar', 0):
OffsetAndMetadata(123, b''),
TopicPartition('foobar', 1):
OffsetAndMetadata(234, b'')
})
coordinator._subscription.assign_from_user([TopicPartition('foobar', 0)])
assert coordinator._subscription.needs_fetch_committed_offsets is True
coordinator.refresh_committed_offsets_if_needed()
assignment = coordinator._subscription.assignment
assert assignment[TopicPartition('foobar', 0)].committed == 123
assert TopicPartition('foobar', 1) not in assignment
assert coordinator._subscription.needs_fetch_committed_offsets is False
def _list_consumer_group_offsets_process_response(self, response):
"""Process an OffsetFetchResponse.
:param response: an OffsetFetchResponse.
:return: A dictionary composed of TopicPartition keys and
OffsetAndMetada values.
"""
if response.API_VERSION <= 3:
# OffsetFetchResponse_v1 lacks a top-level error_code
if response.API_VERSION > 1:
error_type = Errors.for_code(response.error_code)
if error_type is not Errors.NoError:
# optionally we could retry if error_type.retriable
raise error_type(
"OffsetFetchResponse failed with response '{}'."
.format(response))
# transform response into a dictionary with TopicPartition keys and
# OffsetAndMetada values--this is what the Java AdminClient returns
offsets = {}
for topic, partitions in response.topics:
for partition, offset, metadata, error_code in partitions:
error_type = Errors.for_code(error_code)
if error_type is not Errors.NoError:
raise error_type(
"Unable to fetch consumer group offsets for topic {}, partition {}"
.format(topic, partition))
offsets[TopicPartition(topic, partition)] = OffsetAndMetadata(offset, metadata)
else:
raise NotImplementedError(
"Support for OffsetFetchResponse_v{} has not yet been added to KafkaAdminClient."
.format(response.API_VERSION))
return offsets
def main(server, topic, part, port):
server_port = "{}:{}".format(server,port)
consumer = KafkaConsumer(bootstrap_servers=[server_port],
auto_offset_reset='earliest',
enable_auto_commit=False,
value_deserializer=lambda x: json.loads(x.decode('utf-8')))
consumer.assign([TopicPartition(topic, part)])
events = list()
distutils.dir_util.mkpath(os.path.join("data", server)
try:
for event in consumer:
events.append({
"topic": event.topic,
"partition": event.partition,
"offset": event.offset,
"key": event.key,
"value": event.value,
"headers": {str(k):str(v) for k,v in event.headers},
"checksum": event.checksum
})
if event.offset % 100 == 0:
print(event.offset)
except KeyboardInterrupt:
print("Saving Data")
with open('{}/{}_{}.json'.format(server,topic, part), 'w') as _file:
json.dump(events, _file, indent=4)
if __name__ == "__main__":
main()
async def seek_to(self, internal_name: Tuple[str, str], content_type: str,
request_data: dict):
self.log.info("Resetting offsets for %s to %r", internal_name,
request_data)
self._assert_consumer_exists(internal_name, content_type)
self._assert_has_key(request_data, "offsets", content_type)
seeks = []
for el in request_data["offsets"]:
self._assert_has_key(el, "topic", content_type)
for k in ["offset", "partition"]:
self._assert_has_key(el, k, content_type)
convert_to_int(el, k, content_type)
self._assert_positive_number(el, "offset", content_type)
seeks.append(
(TopicPartition(topic=el["topic"],
partition=el["partition"]), el["offset"]))
async with self.consumer_locks[internal_name]:
consumer = self.consumers[internal_name].consumer
for part, offset in seeks:
try:
consumer.seek(part, offset)
except AssertionError:
self._illegal_state_fail(f"Partition {part} is unassigned",
content_type)
empty_response()
async def seek_limit(self,
internal_name: Tuple[str, str],
content_type: str,
request_data: dict,
beginning: bool = True):
direction = "beginning" if beginning else "end"
self.log.info("Seeking %s offsets", direction)
self._assert_consumer_exists(internal_name, content_type)
self._assert_has_key(request_data, "partitions", content_type)
resets = []
for el in request_data["partitions"]:
convert_to_int(el, "partition", content_type)
for k in ["topic", "partition"]:
self._assert_has_key(el, k, content_type)
resets.append(
TopicPartition(topic=el["topic"], partition=el["partition"]))
async with self.consumer_locks[internal_name]:
consumer = self.consumers[internal_name].consumer
try:
if beginning:
consumer.seek_to_beginning(*resets)
else:
consumer.seek_to_end(*resets)
empty_response()
except AssertionError:
self._illegal_state_fail(
f"Trying to reset unassigned partitions to {direction}",
content_type)
def test_read_groups_with_partition(self):
kafka_config = mock.Mock()
kafka_group_reader = KafkaGroupReader(kafka_config)
with mock.patch(
'kafka_utils.kafka_consumer_manager.util.KafkaConsumer',
) as mock_consumer:
with mock.patch.object(
kafka_group_reader,
'get_current_watermarks',
return_value={0: PartitionOffsets('test_topic', 0, 45, 0)},
autospec=True):
with mock.patch.object(
kafka_group_reader,
'parse_consumer_offset_message',
return_value=['test_group', 'test_topic', 0, 45],
autospec=True):
mock_consumer.return_value.__iter__.return_value = iter([
mock.Mock(partition=0, topic='test_topic', offset=45)
])
kafka_group_reader.read_groups(partition=0)
assert kafka_group_reader.kafka_groups['test_group'] == {
"test_topic"
}
mock_consumer.return_value.assign.assert_called_once_with(
[TopicPartition("__consumer_offsets", 0)])
async def commit_offsets(self, internal_name: Tuple[str, str],
content_type: str, request_data: dict,
cluster_metadata: dict):
self.log.info("Committing offsets for %s", internal_name)
self._assert_consumer_exists(internal_name, content_type)
if request_data:
self._assert_has_key(request_data, "offsets", content_type)
payload = {}
for el in request_data.get("offsets", []):
for k in ["partition", "offset"]:
convert_to_int(el, k, content_type)
# If we commit for a partition that does not belong to this consumer, then the internal error raised
# is marked as retriable, and thus the commit method will remain blocked in what looks like an infinite loop
self._topic_and_partition_valid(cluster_metadata, el, content_type)
payload[TopicPartition(el["topic"],
el["partition"])] = OffsetAndMetadata(
el["offset"] + 1, None)
async with self.consumer_locks[internal_name]:
consumer = self.consumers[internal_name].consumer
payload = payload or None
try:
consumer.commit(offsets=payload)
except KafkaError as e:
KarapaceBase.internal_error(
message=f"error sending commit request: {e}",
content_type=content_type)
empty_response()
def notification():
try:
login_flag = session['loggedin_user']
except Exception:
error = "No Logged in User"
return error
username = session['username_user']
heading = "Notification for username: "******"<br><br>"
filepathkafka = "configuration/kafka_config.json"
kafka_ip, kafka_port = read_json_kafka(filepathkafka)
# '{}:{}'.format(kafka_ip,kafka_port)
consumer = KafkaConsumer(bootstrap_servers = ['{}:{}'.format(kafka_ip,kafka_port)], api_version = (0,10),consumer_timeout_ms = 1000, auto_offset_reset="earliest")
# consumer.seek_to_beginning()
partition = TopicPartition(username, 0)
consumer.assign([partition])
consumer.seek(partition, 0)
notification_string = ""
for message in consumer:
notification_string += str(message.value) + "<br>"
# notification_string += "<br><br><br><br>" + str(message.offset)
return heading + notification_string + all_link_str
def rotate_partition(self, partition):
if int(self.partitions[partition].offset) == int(
self.partitions[partition].header.get_start_offset()):
self.logger.debug("Skiping rotate for partition " + partition +
". No new writes")
return
self.logger.info("I need to rotate " + partition)
self.partitions[partition].writer.close()
start_offset = self.partitions[partition].header.get_start_offset()
end_offset = self.partitions[partition].offset
topic = self.partitions[partition].header.get_topic()
part_number = self.partitions[partition].header.get_partition()
dest_dir = os.path.join(self.config.get("main", "working_directory"),
"tocompress", topic, str(part_number))
date = time.strftime("%y%m%d")
dest_filename = os.path.join(
dest_dir, topic + "-" + str(part_number) + "_" +
str(start_offset) + "-" + str(end_offset) + "_" + date + ".mak")
self.mkdirp(dest_dir)
os.rename(self.partitions[partition].writer.get_filename(),
dest_filename)
self.create_new_data_file(
self.partitions[partition].writer.get_filename(),
TopicPartition(topic=topic, partition=part_number))
def fetcher(client, subscription_state):
subscription_state.subscribe(topics=['foobar'])
assignment = [TopicPartition('foobar', i) for i in range(3)]
subscription_state.assign_from_subscribed(assignment)
for tp in assignment:
subscription_state.seek(tp, 0)
return Fetcher(client, subscription_state, Metrics())
def _handle_produce_response(self, node_id, send_time, batches, response):
"""Handle a produce response."""
# if we have a response, parse it
log.debug('Parsing produce response: %r', response)
if response:
batches_by_partition = dict([(batch.topic_partition, batch)
for batch in batches])
for topic, partitions in response.topics:
for partition_info in partitions:
if response.API_VERSION < 2:
partition, error_code, offset = partition_info
ts = None
else:
partition, error_code, offset, ts = partition_info
tp = TopicPartition(topic, partition)
error = Errors.for_code(error_code)
batch = batches_by_partition[tp]
self._complete_batch(batch, error, offset, ts)
if response.API_VERSION > 0:
self._sensors.record_throttle_time(response.throttle_time_ms,
node=node_id)
else:
# this is the acks = 0 case, just complete all requests
for batch in batches:
self._complete_batch(batch, None, -1, None)
def _handle_offset_response(self, future, response):
"""Callback for the response of the list offset call above.
Arguments:
future (Future): the future to update based on response
response (OffsetResponse): response from the server
Raises:
AssertionError: if response does not match partition
"""
timestamp_offset_map = {}
for topic, part_data in response.topics:
for partition_info in part_data:
partition, error_code = partition_info[:2]
partition = TopicPartition(topic, partition)
error_type = Errors.for_code(error_code)
if error_type is Errors.NoError:
if response.API_VERSION == 0:
offsets = partition_info[2]
assert len(offsets) <= 1, 'Expected OffsetResponse with one offset'
if not offsets:
offset = UNKNOWN_OFFSET
else:
offset = offsets[0]
log.debug("Handling v0 ListOffsetResponse response for %s. "
"Fetched offset %s", partition, offset)
if offset != UNKNOWN_OFFSET:
timestamp_offset_map[partition] = (offset, None)
else:
timestamp, offset = partition_info[2:]
log.debug("Handling ListOffsetResponse response for %s. "
"Fetched offset %s, timestamp %s",
partition, offset, timestamp)
if offset != UNKNOWN_OFFSET:
timestamp_offset_map[partition] = (offset, timestamp)
elif error_type is Errors.UnsupportedForMessageFormatError:
# The message format on the broker side is before 0.10.0,
# we simply put None in the response.
log.debug("Cannot search by timestamp for partition %s because the"
" message format version is before 0.10.0", partition)
elif error_type is Errors.NotLeaderForPartitionError:
log.debug("Attempt to fetch offsets for partition %s failed due"
" to obsolete leadership information, retrying.",
partition)
future.failure(error_type(partition))
return
elif error_type is Errors.UnknownTopicOrPartitionError:
log.warning("Received unknown topic or partition error in ListOffset "
"request for partition %s. The topic/partition " +
"may not exist or the user may not have Describe access "
"to it.", partition)
future.failure(error_type(partition))
return
else:
log.warning("Attempt to fetch offsets for partition %s failed due to:"
" %s", partition, error_type)
future.failure(error_type(partition))
return
if not future.is_done:
future.success(timestamp_offset_map)
def test_get_leader_for_partitions_reloads_metadata(self, protocol, conn):
"Get leader for partitions reload metadata if it is not available"
mock_conn(conn)
brokers = [
BrokerMetadata(0, 'broker_1', 4567, None),
BrokerMetadata(1, 'broker_2', 5678, None)
]
resp0_brokers = list(map(itemgetter(0, 1, 2), brokers))
topics = [(NO_LEADER, 'topic_no_partitions', [])]
protocol.decode_metadata_response.return_value = MetadataResponse[0](
resp0_brokers, topics)
client = SimpleClient(hosts=['broker_1:4567'])
# topic metadata is loaded but empty
self.assertDictEqual({}, client.topics_to_brokers)
topics = [(NO_ERROR, 'topic_one_partition', [(NO_ERROR, 0, 0, [0, 1],
[0, 1])])]
protocol.decode_metadata_response.return_value = MetadataResponse[0](
resp0_brokers, topics)
# calling _get_leader_for_partition (from any broker aware request)
# will try loading metadata again for the same topic
leader = client._get_leader_for_partition('topic_one_partition', 0)
self.assertEqual(brokers[0], leader)
self.assertDictEqual(
{TopicPartition('topic_one_partition', 0): brokers[0]},
client.topics_to_brokers)
def test_read_groups(self):
kafka_config = mock.Mock()
kafka_group_reader = KafkaGroupReader(kafka_config)
with mock.patch(
'kafka_utils.kafka_consumer_manager.util.KafkaConsumer',
autospec=True) as mock_consumer:
with mock.patch.object(
kafka_group_reader,
'get_current_watermarks',
return_value={
0: PartitionOffsets('__consumer_offsets', 0, 45, 0),
1: PartitionOffsets('__consumer_offsets', 1, 20, 0),
2: PartitionOffsets('__consumer_offsets', 2, 25, 25),
3: PartitionOffsets('__consumer_offsets', 3, 0, 0),
},
autospec=True):
with mock.patch.object(
kafka_group_reader,
'parse_consumer_offset_message',
side_effect=iter([
('test_group', 'test_topic', 0, 45),
('test_group2', 'test_topic2', 0, 20),
]),
autospec=True,
):
mock_consumer.return_value.next.side_effect = iter([
mock.Mock(offset=44, partition=0, topic='test_topic'),
mock.Mock(offset=19, partition=1, topic='test_topic'),
])
mock_consumer.return_value.partitions_for_topic.return_value = [
0, 1
]
kafka_group_reader.read_groups()
assert kafka_group_reader.kafka_groups['test_group'] == {
"test_topic"
}
assert kafka_group_reader.kafka_groups['test_group2'] == {
"test_topic2"
}
mock_consumer.return_value.assign.call_args_list == [
mock.call([
TopicPartition("__consumer_offsets", 0),
TopicPartition("__consumer_offsets", 1),
]),
mock.call([TopicPartition("__consumer_offsets", 0)]),
]
def test_normal_consume(self, KafkaConsumer):
# Make a fake consumer to generate a message
fake_kafka_consumer = self.mock_consumer(
KafkaConsumer,
value=b'json:{"message":{"code":"NY","name":"New York"},"version":1,"type":"us-state"}',
max_calls=100,
)
# Test the values sent to our serializer match the message
def save(ser):
self.assertEqual(ser.validated_data["code"], "NY")
self.assertEqual(ser.validated_data["name"], "New York")
FakeStateSerializer = self.mock_state_serializer(save)
# Consume a message
consumer = Consumer(TOPIC_STATES, consumer_timeout_ms=500)
consumer.register(FakeStateSerializer)
consumer.run(iter_limit=1)
# Test the expected mocks where called
KafkaConsumer.assert_called_once_with(
auto_offset_reset="earliest",
bootstrap_servers=["kafka:9092"],
consumer_timeout_ms=500,
enable_auto_commit=False,
)
fake_kafka_consumer.partitions_for_topic.assert_called_once_with(TOPIC_STATES)
fake_kafka_consumer.assign.assert_called_once_with(
[
TopicPartition(partition=0, topic=TOPIC_STATES),
TopicPartition(partition=1, topic=TOPIC_STATES),
]
)
self.assertEqual(KafkaConsumer.call_count, 1)
self.assertEqual(FakeStateSerializer.call_count, 1)
self.assertEqual(fake_kafka_consumer.__next__.call_count, 1)
self.assertEqual(self.serializers["state"].save.call_count, 1)
consumer.run(iter_limit=1)
self.assertEqual(KafkaConsumer.call_count, 1)
self.assertEqual(FakeStateSerializer.call_count, 2)
self.assertEqual(fake_kafka_consumer.__next__.call_count, 2)
self.assertEqual(self.serializers["state"].save.call_count, 1)
def commit_offsets_in_kafka(cons, topic, offsets_topic):
for partition, offset in offsets_topic.items():
tp = TopicPartition(topic, int(partition))
cons.assign([tp])
cons.seek(tp, int(offset))
# commit it
cons.commit()
def get_assignments(consumer, *topics):
assignments = []
for topic in topics:
partitions = consumer.partitions_for_topic(topic)
for p in partitions:
assignments.append(TopicPartition(topic, p))
return assignments
def test_subscription_listener(mocker, coordinator):
listener = mocker.MagicMock(spec=ConsumerRebalanceListener)
coordinator._subscription.subscribe(
topics=['foobar'],
listener=listener)
coordinator._on_join_prepare(0, 'member-foo')
assert listener.on_partitions_revoked.call_count == 1
listener.on_partitions_revoked.assert_called_with(set([]))
assignment = ConsumerProtocolMemberAssignment(0, [('foobar', [0, 1])], b'')
coordinator._on_join_complete(
0, 'member-foo', 'roundrobin', assignment.encode())
assert listener.on_partitions_assigned.call_count == 1
listener.on_partitions_assigned.assert_called_with(set([
TopicPartition('foobar', 0),
TopicPartition('foobar', 1)]))
def test_get_leader_exceptions_when_noleader(self, protocol, conn):
mock_conn(conn)
brokers = [
BrokerMetadata(0, 'broker_1', 4567, None),
BrokerMetadata(1, 'broker_2', 5678, None)
]
resp0_brokers = list(map(itemgetter(0, 1, 2), brokers))
topics = [
(NO_ERROR, 'topic_noleader', [
(NO_LEADER, 0, -1, [], []),
(NO_LEADER, 1, -1, [], []),
]),
]
protocol.decode_metadata_response.return_value = MetadataResponse[0](resp0_brokers, topics)
client = SimpleClient(hosts=['broker_1:4567'])
self.assertDictEqual(
{
TopicPartition('topic_noleader', 0): None,
TopicPartition('topic_noleader', 1): None
},
client.topics_to_brokers)
# No leader partitions -- raise LeaderNotAvailableError
with self.assertRaises(LeaderNotAvailableError):
self.assertIsNone(client._get_leader_for_partition('topic_noleader', 0))
with self.assertRaises(LeaderNotAvailableError):
self.assertIsNone(client._get_leader_for_partition('topic_noleader', 1))
# Unknown partitions -- raise UnknownTopicOrPartitionError
with self.assertRaises(UnknownTopicOrPartitionError):
self.assertIsNone(client._get_leader_for_partition('topic_noleader', 2))
topics = [
(NO_ERROR, 'topic_noleader', [
(NO_ERROR, 0, 0, [0, 1], [0, 1]),
(NO_ERROR, 1, 1, [1, 0], [1, 0])
]),
]
protocol.decode_metadata_response.return_value = MetadataResponse[0](resp0_brokers, topics)
self.assertEqual(brokers[0], client._get_leader_for_partition('topic_noleader', 0))
self.assertEqual(brokers[1], client._get_leader_for_partition('topic_noleader', 1))
def test__parse_fetched_data__out_of_range(fetcher, topic, mocker):
fetcher.config['check_crcs'] = False
tp = TopicPartition(topic, 0)
completed_fetch = CompletedFetch(tp, 0, 0,
[OffsetOutOfRangeError.errno, -1, None],
mocker.MagicMock())
partition_record = fetcher._parse_fetched_data(completed_fetch)
assert partition_record is None
assert fetcher._subscriptions.assignment[tp].awaiting_reset is True
def test__parse_fetched_data__unknown_tp(fetcher, topic, mocker):
fetcher.config['check_crcs'] = False
tp = TopicPartition(topic, 0)
completed_fetch = CompletedFetch(
tp, 0, 0, [UnknownTopicOrPartitionError.errno, -1, None],
mocker.MagicMock())
partition_record = fetcher._parse_fetched_data(completed_fetch)
assert partition_record is None
fetcher._client.cluster.request_update.assert_called_with()
