예제 #1
0
    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()
예제 #2
0
    def _send_offset_request(self, partition, timestamp):
        """Fetch a single offset before the given timestamp for the partition.

        Arguments:
            partition (TopicPartition): partition that needs fetching offset
            timestamp (int): timestamp for fetching offset

        Returns:
            Future: resolves to the corresponding offset
        """
        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))

        request = OffsetRequest[0](
            -1, [(partition.topic, [(partition.partition, timestamp, 1)])]
        )
        # 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
        future = Future()
        _f = self._client.send(node_id, request)
        _f.add_callback(self._handle_offset_response, partition, future)
        _f.add_errback(lambda e: future.failure(e))
        return future
예제 #3
0
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)
예제 #4
0
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) 
예제 #5
0
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
예제 #6
0
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)
예제 #7
0
파일: conn.py 프로젝트: Yelp/kafka-python
    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)
예제 #8
0
    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)
예제 #9
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()
예제 #10
0
 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)
예제 #11
0
    def _send(self, request, expect_response=True):
        assert self.state in (ConnectionStates.AUTHENTICATING, ConnectionStates.CONNECTED)
        future = Future()
        correlation_id = self._next_correlation_id()
        header = RequestHeader(request,
                               correlation_id=correlation_id,
                               client_id=self.config['client_id'])
        message = b''.join([header.encode(), request.encode()])
        size = Int32.encode(len(message))
        data = size + message
        try:
            # In the future we might manage an internal write buffer
            # and send bytes asynchronously. For now, just block
            # sending each request payload
            self._sock.setblocking(True)
            total_sent = 0
            while total_sent < len(data):
                sent_bytes = self._sock.send(data[total_sent:])
                total_sent += sent_bytes
            assert total_sent == len(data)
            if self._sensors:
                self._sensors.bytes_sent.record(total_sent)
            self._sock.setblocking(False)
        except (AssertionError, ConnectionError) as e:
            log.exception("Error sending %s to %s", request, self)
            error = Errors.ConnectionError("%s: %s" % (str(self), e))
            self.close(error=error)
            return future.failure(error)
        log.debug('%s Request %d: %s', self, correlation_id, request)

        if expect_response:
            ifr = InFlightRequest(request=request,
                                  correlation_id=correlation_id,
                                  response_type=request.RESPONSE_TYPE,
                                  future=future,
                                  timestamp=time.time())
            self.in_flight_requests.append(ifr)
        else:
            future.success(None)

        return future
예제 #12
0
 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
예제 #13
0
파일: conn.py 프로젝트: Yelp/kafka-python
    def _send(self, request, expect_response=True):
        future = Future()
        correlation_id = self._next_correlation_id()
        header = RequestHeader(request,
                               correlation_id=correlation_id,
                               client_id=self.config['client_id'])
        message = b''.join([header.encode(), request.encode()])
        size = Int32.encode(len(message))
        data = size + message
        try:
            # In the future we might manage an internal write buffer
            # and send bytes asynchronously. For now, just block
            # sending each request payload
            self._sock.setblocking(True)
            total_sent = 0
            while total_sent < len(data):
                sent_bytes = self._sock.send(data[total_sent:])
                total_sent += sent_bytes
            assert total_sent == len(data)
            if self._sensors:
                self._sensors.bytes_sent.record(total_sent)
            self._sock.setblocking(False)
        except (AssertionError, ConnectionError) as e:
            log.exception("Error sending %s to %s", request, self)
            error = Errors.ConnectionError("%s: %s" % (str(self), e))
            self.close(error=error)
            return future.failure(error)
        log.debug('%s Request %d: %s', self, correlation_id, request)

        if expect_response:
            ifr = InFlightRequest(request=request,
                                  correlation_id=correlation_id,
                                  response_type=request.RESPONSE_TYPE,
                                  future=future,
                                  timestamp=time.time())
            self.in_flight_requests.append(ifr)
        else:
            future.success(None)

        return future
예제 #14
0
    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)
            _f.add_errback(lambda e: future_request.failure(e))
            futures.append(future_request)
        return futures
예제 #15
0
    def test_send_broker_unaware_request(self):
        mocked_conns = {("kafka01", 9092): MagicMock(), ("kafka02", 9092): MagicMock(), ("kafka03", 9092): MagicMock()}
        # inject BrokerConnection 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, afi):
            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()
예제 #16
0
    def _send_offset_request(self, node_id, timestamps):
        by_topic = collections.defaultdict(list)
        for tp, timestamp in six.iteritems(timestamps):
            if self.config['api_version'] >= (0, 10, 1):
                data = (tp.partition, timestamp)
            else:
                data = (tp.partition, timestamp, 1)
            by_topic[tp.topic].append(data)

        if self.config['api_version'] >= (0, 10, 1):
            request = OffsetRequest[1](-1, list(six.iteritems(by_topic)))
        else:
            request = OffsetRequest[0](-1, list(six.iteritems(by_topic)))

        # 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
        future = Future()

        _f = self._client.send(node_id, request)
        _f.add_callback(self._handle_offset_response, future)
        _f.add_errback(lambda e: future.failure(e))
        return future
예제 #17
0
파일: conn.py 프로젝트: arpras/kafka-python
    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)))
        correlation_id = self._next_correlation_id()
        header = RequestHeader(request,
                               correlation_id=correlation_id,
                               client_id=self.config['client_id'])
        message = b''.join([header.encode(), request.encode()])
        size = Int32.encode(len(message))
        try:
            # In the future we might manage an internal write buffer
            # and send bytes asynchronously. For now, just block
            # sending each request payload
            self._sock.setblocking(True)
            for data in (size, message):
                total_sent = 0
                while total_sent < len(data):
                    sent_bytes = self._sock.send(data[total_sent:])
                    total_sent += sent_bytes
                assert total_sent == len(data)
            self._sock.setblocking(False)
        except (AssertionError, ConnectionError) as e:
            log.exception("Error sending %s to %s", request, self)
            error = Errors.ConnectionError("%s: %s" % (str(self), e))
            self.close(error=error)
            return future.failure(error)
        log.debug('%s Request %d: %s', self, correlation_id, request)

        if expect_response:
            ifr = InFlightRequest(request=request,
                                  correlation_id=correlation_id,
                                  response_type=request.RESPONSE_TYPE,
                                  future=future,
                                  timestamp=time.time())
            self.in_flight_requests.append(ifr)
        else:
            future.success(None)

        return future
예제 #18
0
 def test_delete_topics(
     self,
     mock_new_topics,
     mock_least_loaded_node,
     bootstrap_brokers,
     delete_response,
     metadata_response,
 ):
     mock_kafka_client = mock.Mock()
     mock_kafka_client.poll = \
         mock.Mock(side_effect=[metadata_response, delete_response])
     mock_kafka_client.ready.return_value = True
     mock_kafka_client.least_loaded_node.return_value = \
         mock_least_loaded_node
     mock_kafka_client.send.return_value = Future()
     admin = AdminClient(mock_kafka_client)
     response = admin.delete_topics(mock_new_topics, 0)
     assert response == delete_response
예제 #19
0
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)
예제 #20
0
    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)

        version = 0 if self.config["api_version"] < (0, 11, 0) else 1
        request = SyncGroupRequest[version](
            self.group_id,
            self._generation.generation_id,
            self._generation.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)