def test_handle_offset_commit_response(mocker, patched_coord, offsets, response, error, dead): future = Future() patched_coord._handle_offset_commit_response(offsets, future, time.time(), response) assert isinstance(future.exception, error) assert patched_coord.coordinator_id is (None if dead else 0)
def send(self, node_id, request, wakeup=True): """Send a request to a specific node. Bytes are placed on an internal per-connection send-queue. Actual network I/O will be triggered in a subsequent call to .poll() Arguments: node_id (int): destination node request (Struct): request object (not-encoded) wakeup (bool): optional flag to disable thread-wakeup Raises: AssertionError: if node_id is not in current cluster metadata Returns: Future: resolves to Response struct or Error """ conn = self._conns.get(node_id) if not conn or not self._can_send_request(node_id): self.maybe_connect(node_id, wakeup=wakeup) return Future().failure(Errors.NodeNotReadyError(node_id)) # conn.send will queue the request internally # we will need to call send_pending_requests() # to trigger network I/O future = conn.send(request, blocking=False) self._sending.add(conn) # Wakeup signal is useful in case another thread is # blocked waiting for incoming network traffic while holding # the client lock in poll(). if wakeup: self.wakeup() return future
def _on_join_leader(self, response): """ Perform leader synchronization and send back the assignment for the group via SyncGroupRequest Arguments: response (JoinResponse): broker response to parse Returns: Future: resolves to member assignment encoded-bytes """ try: group_assignment = self._perform_assignment(response.leader_id, response.group_protocol, response.members) except Exception as e: return Future().failure(e) request = SyncGroupRequest[0]( self.group_id, self.generation, self.member_id, [(member_id, assignment if isinstance(assignment, bytes) else assignment.encode()) for member_id, assignment in six.iteritems(group_assignment)]) log.debug("Sending leader SyncGroup for group %s to coordinator %s: %s", self.group_id, self.coordinator_id, request) return self._send_sync_group_request(request)
def conn(mocker): """Return a connection mocker fixture""" from kafka.conn import ConnectionStates from kafka.future import Future from kafka.protocol.metadata import MetadataResponse conn = mocker.patch('kafka.client_async.BrokerConnection') conn.return_value = conn conn.state = ConnectionStates.CONNECTED conn.send.return_value = Future().success(MetadataResponse[0]( [(0, 'foo', 12), (1, 'bar', 34)], # brokers [])) # topics conn.blacked_out.return_value = False def _set_conn_state(state): conn.state = state return state conn._set_conn_state = _set_conn_state conn.connect.side_effect = lambda: conn.state conn.connect_blocking.return_value = True conn.connecting = lambda: conn.state in (ConnectionStates.CONNECTING, ConnectionStates.HANDSHAKE) conn.connected = lambda: conn.state is ConnectionStates.CONNECTED conn.disconnected = lambda: conn.state is ConnectionStates.DISCONNECTED return conn
def test_send_broker_unaware_request(self): mocked_conns = { ('kafka01', 9092): MagicMock(), ('kafka02', 9092): MagicMock(), ('kafka03', 9092): MagicMock() } # inject KafkaConnection side effects mock_conn(mocked_conns[('kafka01', 9092)], success=False) mock_conn(mocked_conns[('kafka03', 9092)], success=False) future = Future() mocked_conns[('kafka02', 9092)].send.return_value = future mocked_conns[( 'kafka02', 9092)].recv.side_effect = lambda: future.success('valid response') def mock_get_conn(host, port): return mocked_conns[(host, port)] # patch to avoid making requests before we want it with patch.object(SimpleClient, 'load_metadata_for_topics'): with patch.object(SimpleClient, '_get_conn', side_effect=mock_get_conn): client = SimpleClient(hosts='kafka01:9092,kafka02:9092') resp = client._send_broker_unaware_request( payloads=['fake request'], encoder_fn=MagicMock(), decoder_fn=lambda x: x) self.assertEqual('valid response', resp) mocked_conns[('kafka02', 9092)].recv.assert_called_once_with()
def test_lookup_coordinator_failure(mocker, coordinator): mocker.patch.object(coordinator, '_send_group_coordinator_request', return_value=Future().failure(Exception('foobar'))) future = coordinator.lookup_coordinator() assert future.failed()
def test_handle_offset_commit_response(patched_coord, offsets, response, error, dead, reassign): future = Future() patched_coord._handle_offset_commit_response(offsets, future, response) assert isinstance(future.exception, error) assert patched_coord.coordinator_id is (None if dead else 0) assert patched_coord._subscription.needs_partition_assignment is reassign
def _send_offset_requests(self, timestamps): """Fetch offsets for each partition in timestamps dict. This may send request to multiple nodes, based on who is Leader for partition. Arguments: timestamps (dict): {TopicPartition: int} mapping of fetching timestamps. Returns: Future: resolves to a mapping of retrieved offsets """ timestamps_by_node = collections.defaultdict(dict) for partition, timestamp in six.iteritems(timestamps): node_id = self._client.cluster.leader_for_partition(partition) if node_id is None: self._client.add_topic(partition.topic) log.debug("Partition %s is unknown for fetching offset," " wait for metadata refresh", partition) return Future().failure(Errors.StaleMetadata(partition)) elif node_id == -1: log.debug("Leader for partition %s unavailable for fetching " "offset, wait for metadata refresh", partition) return Future().failure( Errors.LeaderNotAvailableError(partition)) else: timestamps_by_node[node_id][partition] = timestamp # Aggregate results until we have all list_offsets_future = Future() responses = [] node_count = len(timestamps_by_node) def on_success(value): responses.append(value) if len(responses) == node_count: offsets = {} for r in responses: offsets.update(r) list_offsets_future.success(offsets) def on_fail(err): if not list_offsets_future.is_done: list_offsets_future.failure(err) for node_id, timestamps in six.iteritems(timestamps_by_node): _f = self._send_offset_request(node_id, timestamps) _f.add_callback(on_success) _f.add_errback(on_fail) return list_offsets_future
def test_commit_offsets_async(mocker, coordinator, offsets): mocker.patch.object(coordinator._client, 'poll') mocker.patch.object(coordinator, 'ensure_coordinator_known') mocker.patch.object(coordinator, '_send_offset_commit_request', return_value=Future().success('fizzbuzz')) ret = coordinator.commit_offsets_async(offsets) assert isinstance(ret, Future) assert coordinator._send_offset_commit_request.call_count == 1
def test_commit_offsets_async(mocker, coordinator, offsets): mocker.patch.object(coordinator._client, 'poll') mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False) mocker.patch.object(coordinator, 'ensure_coordinator_ready') mocker.patch.object(coordinator, '_send_offset_commit_request', return_value=Future().success('fizzbuzz')) coordinator.commit_offsets_async(offsets) assert coordinator._send_offset_commit_request.call_count == 1
def _send_group_metadata_request(self): """Discover the current coordinator for the group. Returns: Future: resolves to the node id of the coordinator """ node_id = self._client.least_loaded_node() if node_id is None or not self._client.ready(node_id): return Future().failure(Errors.NoBrokersAvailable()) log.debug("Issuing group metadata request to broker %s", node_id) request = GroupCoordinatorRequest(self.group_id) future = Future() _f = self._client.send(node_id, request) _f.add_callback(self._handle_group_coordinator_response, future) _f.add_errback(self._failed_request, node_id, request, future) return future
def _send_sync_group_request(self, request): if self.coordinator_unknown(): e = Errors.GroupCoordinatorNotAvailableError(self.coordinator_id) return Future().failure(e) # We assume that coordinator is ready if we're sending SyncGroup # as it typically follows a successful JoinGroup # Also note that if client.ready() enforces a metadata priority policy, # we can get into an infinite loop if the leader assignment process # itself requests a metadata update future = Future() _f = self._client.send(self.coordinator_id, request) _f.add_callback(self._handle_sync_group_response, future, time.time()) _f.add_errback(self._failed_request, self.coordinator_id, request, future) return future
def conn(mocker): conn = mocker.patch('kafka.client_async.BrokerConnection') conn.return_value = conn conn.state = ConnectionStates.CONNECTED conn.send.return_value = Future().success( MetadataResponse( [(0, 'foo', 12), (1, 'bar', 34)], # brokers [])) # topics return conn
def test_send_offset_commit_request_success(mocker, patched_coord, offsets): _f = Future() patched_coord._client.send.return_value = _f future = patched_coord._send_offset_commit_request(offsets) (node, request), _ = patched_coord._client.send.call_args response = OffsetCommitResponse[0]([('foobar', [(0, 0), (1, 0)])]) _f.success(response) patched_coord._handle_offset_commit_response.assert_called_with( offsets, future, mocker.ANY, response)
def test_send_offset_fetch_request_success(patched_coord, partitions): _f = Future() patched_coord._client.send.return_value = _f future = patched_coord._send_offset_fetch_request(partitions) (node, request), _ = patched_coord._client.send.call_args response = OffsetFetchResponse[0]([('foobar', [(0, 123, b'', 0), (1, 234, b'', 0)])]) _f.success(response) patched_coord._handle_offset_fetch_response.assert_called_with( future, response)
def _send_offset_request(self, partitions, timestamp): """Fetch a single offset before the given timestamp for the partition. Arguments: partitions iterable of TopicPartition: partitions that needs fetching offset timestamp (int): timestamp for fetching offset Returns: list of Future: resolves to the corresponding offset """ topic = partitions[0].topic nodes_per_partitions = {} for partition in partitions: node_id = self._client.cluster.leader_for_partition(partition) if node_id is None: log.debug("Partition %s is unknown for fetching offset," " wait for metadata refresh", partition) return [Future().failure(Errors.StaleMetadata(partition))] elif node_id == -1: log.debug("Leader for partition %s unavailable for fetching offset," " wait for metadata refresh", partition) return [Future().failure(Errors.LeaderNotAvailableError(partition))] nodes_per_partitions.setdefault(node_id, []).append(partition) # Client returns a future that only fails on network issues # so create a separate future and attach a callback to update it # based on response error codes futures = [] for node_id, partitions in six.iteritems(nodes_per_partitions): request = OffsetRequest[0]( -1, [(topic, [(partition.partition, timestamp, 1) for partition in partitions])] ) future_request = Future() _f = self._client.send(node_id, request) _f.add_callback(self._handle_offset_response, partitions, future_request) def errback(e): log.error("Offset request errback error %s", e) future_request.failure(e) _f.add_errback(errback) futures.append(future_request) return futures
def test_send_offset_fetch_request_failure(patched_coord, partitions): _f = Future() patched_coord._client.send.return_value = _f future = patched_coord._send_offset_fetch_request(partitions) (node, request), _ = patched_coord._client.send.call_args error = Exception() _f.failure(error) patched_coord._failed_request.assert_called_with(0, request, future, error) assert future.failed() assert future.exception is error
def test_process_queue_success_cb(self): """ Verify success callback executed on Future success""" mock_future = Future() mock_future.is_done = True mock_future.value = 1 self.kafka_producer.send = MagicMock(return_value=mock_future) with self.assertRaises(SuccessException) as context: process_queue(self.event_queue, "test", self.kafka_producer, self.success_cb, self.error_cb)
def _send_heartbeat_request(self): """Send a heartbeat request""" request = HeartbeatRequest[0](self.group_id, self.generation, self.member_id) log.debug("Heartbeat: %s[%s] %s", request.group, request.generation_id, request.member_id) #pylint: disable-msg=no-member future = Future() _f = self._client.send(self.coordinator_id, request) _f.add_callback(self._handle_heartbeat_response, future) _f.add_errback(self._failed_request, self.coordinator_id, request, future) return future
def test_handle_offset_fetch_response(patched_coord, offsets, response, error, dead): future = Future() patched_coord._handle_offset_fetch_response(future, response) if error is not None: assert isinstance(future.exception, error) else: assert future.succeeded() assert future.value == offsets assert patched_coord.coordinator_id is (None if dead else 0)
def test_process_queue(self): """ Verify message queue is processed.""" mock_future = Future() mock_future.is_done = True mock_future.value = 1 self.kafka_producer.send = MagicMock(return_value=mock_future) with self.assertRaises(SuccessException) as context: process_queue(self.event_queue, "test", self.kafka_producer, self.success_cb, self.error_cb) # Verify message queue "get" is called self.event_queue.get.assert_called_with()
def test_send_controller_request(self, mock_least_loaded_node, controller_id, bootstrap_brokers, metadata_response): mock_kafka_client = mock.Mock() mock_kafka_client.poll.return_value = metadata_response mock_kafka_client.least_loaded_node.return_value = \ mock_least_loaded_node mock_kafka_client.send.return_value = Future() mock_kafka_client.connected.return_value = True admin = AdminClient(mock_kafka_client) assert admin._send_controller_request() == controller_id
def test_ensure_active_group(mocker, coordinator): coordinator._subscription.subscribe(topics=['foobar']) mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False) mocker.patch.object(coordinator, '_send_join_group_request', return_value=Future().success(True)) mocker.patch.object(coordinator, 'need_rejoin', side_effect=[True, True, False]) mocker.patch.object(coordinator, '_on_join_complete') mocker.patch.object(coordinator, '_heartbeat_thread') coordinator.ensure_active_group() coordinator._send_join_group_request.assert_called_once_with()
def test__handle_offset_response(fetcher, mocker): # Broker returns UnsupportedForMessageFormatError, will omit partition fut = Future() res = OffsetResponse[1]([ ("topic", [(0, 43, -1, -1)]), ("topic", [(1, 0, 1000, 9999)]) ]) fetcher._handle_offset_response(fut, res) assert fut.succeeded() assert fut.value == {TopicPartition("topic", 1): (9999, 1000)} # Broker returns NotLeaderForPartitionError fut = Future() res = OffsetResponse[1]([ ("topic", [(0, 6, -1, -1)]), ]) fetcher._handle_offset_response(fut, res) assert fut.failed() assert isinstance(fut.exception, NotLeaderForPartitionError) # Broker returns UnknownTopicOrPartitionError fut = Future() res = OffsetResponse[1]([ ("topic", [(0, 3, -1, -1)]), ]) fetcher._handle_offset_response(fut, res) assert fut.failed() assert isinstance(fut.exception, UnknownTopicOrPartitionError) # Broker returns many errors and 1 result # Will fail on 1st error and return fut = Future() res = OffsetResponse[1]([ ("topic", [(0, 43, -1, -1)]), ("topic", [(1, 6, -1, -1)]), ("topic", [(2, 3, -1, -1)]), ("topic", [(3, 0, 1000, 9999)]) ]) fetcher._handle_offset_response(fut, res) assert fut.failed() assert isinstance(fut.exception, NotLeaderForPartitionError)
def _send_heartbeat_request(self): """Send a heartbeat request""" if self.coordinator_unknown(): e = Errors.GroupCoordinatorNotAvailableError(self.coordinator_id) return Future().failure(e) elif not self._client.ready(self.coordinator_id, metadata_priority=False): e = Errors.NodeNotReadyError(self.coordinator_id) return Future().failure(e) version = 0 if self.config['api_version'] < (0, 11, 0) else 1 request = HeartbeatRequest[version](self.group_id, self._generation.generation_id, self._generation.member_id) log.debug("Heartbeat: %s[%s] %s", request.group, request.generation_id, request.member_id) # pylint: disable-msg=no-member future = Future() _f = self._client.send(self.coordinator_id, request) _f.add_callback(self._handle_heartbeat_response, future, time.time()) _f.add_errback(self._failed_request, self.coordinator_id, request, future) return future
def send(self, request, expect_response=True): """send request, return Future() Can block on network if request is larger than send_buffer_bytes """ future = Future() if self.connecting(): return future.failure(Errors.NodeNotReadyError(str(self))) elif not self.connected(): return future.failure(Errors.ConnectionError(str(self))) elif not self.can_send_more(): return future.failure(Errors.TooManyInFlightRequests(str(self))) return self._send(request, expect_response=expect_response)
def _send_group_coordinator_request(self): """Discover the current coordinator for the group. Returns: Future: resolves to the node id of the coordinator """ node_id = self._client.least_loaded_node() if node_id is None: return Future().failure(Errors.NoBrokersAvailable()) elif not self._client.ready(node_id, metadata_priority=False): e = Errors.NodeNotReadyError(node_id) return Future().failure(e) log.debug("Sending group coordinator request for group %s to broker %s", self.group_id, node_id) request = GroupCoordinatorRequest[0](self.group_id) future = Future() _f = self._client.send(node_id, request) _f.add_callback(self._handle_group_coordinator_response, future) _f.add_errback(self._failed_request, node_id, request, future) return future
def test_process_queue_error_cb(self): """ Verify error callback executed on Future exception""" mock_future = Future() mock_future.is_done = True mock_future.exception = FailureException() cb_spy = Mock(wraps=self.error_cb) with self.assertRaises(FailureException) as context: self.kafka_producer.send = MagicMock(return_value=mock_future) process_queue(self.event_queue, "test", self.kafka_producer, self.success_cb, cb_spy) # Also verify propagation of queue and message to error callback. # Message is specified by mocking of self.event_queue.get in setUp. cb_spy.assert_called_with(1, self.event_queue)
def test_fetch_committed_offsets(mocker, coordinator): # No partitions, no IO polling mocker.patch.object(coordinator._client, 'poll') assert coordinator.fetch_committed_offsets([]) == {} assert coordinator._client.poll.call_count == 0 # general case -- send offset fetch request, get successful future mocker.patch.object(coordinator, 'ensure_coordinator_known') mocker.patch.object(coordinator, '_send_offset_fetch_request', return_value=Future().success('foobar')) partitions = [TopicPartition('foobar', 0)] ret = coordinator.fetch_committed_offsets(partitions) assert ret == 'foobar' coordinator._send_offset_fetch_request.assert_called_with(partitions) assert coordinator._client.poll.call_count == 1 # Failed future is raised if not retriable coordinator._send_offset_fetch_request.return_value = Future().failure( AssertionError) coordinator._client.poll.reset_mock() try: coordinator.fetch_committed_offsets(partitions) except AssertionError: pass else: assert False, 'Exception not raised when expected' assert coordinator._client.poll.call_count == 1 coordinator._client.poll.reset_mock() coordinator._send_offset_fetch_request.side_effect = [ Future().failure(Errors.RequestTimedOutError), Future().success('fizzbuzz') ] ret = coordinator.fetch_committed_offsets(partitions) assert ret == 'fizzbuzz' assert coordinator._client.poll.call_count == 2 # call + retry
def test_commit_offsets_sync(mocker, coordinator, offsets): mocker.patch.object(coordinator, 'ensure_coordinator_known') mocker.patch.object(coordinator, '_send_offset_commit_request', return_value=Future().success('fizzbuzz')) cli = coordinator._client mocker.patch.object(cli, 'poll') # No offsets, no calls assert coordinator.commit_offsets_sync({}) is None assert coordinator._send_offset_commit_request.call_count == 0 assert cli.poll.call_count == 0 ret = coordinator.commit_offsets_sync(offsets) assert coordinator._send_offset_commit_request.call_count == 1 assert cli.poll.call_count == 1 assert ret == 'fizzbuzz' # Failed future is raised if not retriable coordinator._send_offset_commit_request.return_value = Future().failure( AssertionError) coordinator._client.poll.reset_mock() try: coordinator.commit_offsets_sync(offsets) except AssertionError: pass else: assert False, 'Exception not raised when expected' assert coordinator._client.poll.call_count == 1 coordinator._client.poll.reset_mock() coordinator._send_offset_commit_request.side_effect = [ Future().failure(Errors.RequestTimedOutError), Future().success('fizzbuzz') ] ret = coordinator.commit_offsets_sync(offsets) assert ret == 'fizzbuzz' assert coordinator._client.poll.call_count == 2 # call + retry