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()