コード例 #1
0
    def setUp(self):
        self.config = {
            'host': 'localhost',
            'port': 9090,
            'request_id': 0,
            'payload': b'test data',
            'payload2': b'another packet'
        }

        # Mocking socket.create_connection will cause _sock to always be a
        # MagicMock()
        patcher = mock.patch('socket.create_connection', spec=True)
        self.MockCreateConn = patcher.start()
        self.addCleanup(patcher.stop)

        # Also mock socket.sendall() to appear successful
        self.MockCreateConn().sendall.return_value = None

        # And mock socket.recv() to return two payloads, then '', then raise
        # Note that this currently ignores the num_bytes parameter to sock.recv()
        payload_size = len(self.config['payload'])
        payload2_size = len(self.config['payload2'])
        self.MockCreateConn().recv.side_effect = [
            struct.pack('>i', payload_size),
            struct.pack('>%ds' % payload_size, self.config['payload']),
            struct.pack('>i', payload2_size),
            struct.pack('>%ds' % payload2_size, self.config['payload2']), b''
        ]

        # Create a connection object
        self.conn = KafkaConnection(self.config['host'], self.config['port'])

        # Reset any mock counts caused by __init__
        self.MockCreateConn.reset_mock()
コード例 #2
0
ファイル: test_conn.py プロジェクト: yobibytes/kafka-python
    def test_copy(self, socket):
        """KafkaConnection copies work as expected"""

        conn = KafkaConnection('kafka', 9092)
        self.assertEqual(socket.call_count, 1)

        copy = conn.copy()
        self.assertEqual(socket.call_count, 1)
        self.assertEqual(copy.host, 'kafka')
        self.assertEqual(copy.port, 9092)
        self.assertEqual(copy._sock, None)

        copy.reinit()
        self.assertEqual(socket.call_count, 2)
        self.assertNotEqual(copy._sock, None)
コード例 #3
0
ファイル: test_conn.py プロジェクト: AaronSense/kafka-python
    def setUp(self):
        self.config = {
            'host': 'localhost',
            'port': 9090,
            'request_id': 0,
            'payload': b'test data',
            'payload2': b'another packet'
        }

        # Mocking socket.create_connection will cause _sock to always be a
        # MagicMock()
        patcher = mock.patch('socket.create_connection', spec=True)
        self.MockCreateConn = patcher.start()
        self.addCleanup(patcher.stop)

        # Also mock socket.sendall() to appear successful
        self.MockCreateConn().sendall.return_value = None

        # And mock socket.recv() to return two payloads, then '', then raise
        # Note that this currently ignores the num_bytes parameter to sock.recv()
        payload_size = len(self.config['payload'])
        payload2_size = len(self.config['payload2'])
        self.MockCreateConn().recv.side_effect = [
            struct.pack('>i', payload_size),
            struct.pack('>%ds' % payload_size, self.config['payload']),
            struct.pack('>i', payload2_size),
            struct.pack('>%ds' % payload2_size, self.config['payload2']),
            b''
        ]

        # Create a connection object
        self.conn = KafkaConnection(self.config['host'], self.config['port'])

        # Reset any mock counts caused by __init__
        self.MockCreateConn.reset_mock()
コード例 #4
0
    def _get_conn_for_broker(self, broker):
        "Get or create a connection to a broker"
        if (broker.host, broker.port) not in self.conns:
            self.conns[(broker.host, broker.port)] = \
                KafkaConnection(broker.host, broker.port, self.bufsize)

        return self.conns[(broker.host, broker.port)]
コード例 #5
0
ファイル: client.py プロジェクト: secat/kafka-python
    def _get_conn(self, host, port):
        "Get or create a connection to a broker using host and port"

        host_key = (host, port)
        if host_key not in self.conns:
            self.conns[host_key] = KafkaConnection(host, port)

        return self.conns[host_key]
コード例 #6
0
    def test_init_failure_raises_connection_error(self):
        def raise_error(*args):
            raise socket.error

        assert socket.create_connection is self.MockCreateConn
        socket.create_connection.side_effect = raise_error
        with self.assertRaises(ConnectionError):
            KafkaConnection(self.config['host'], self.config['port'])
コード例 #7
0
ファイル: client.py プロジェクト: secat/kafka-python
    def _get_conn_for_broker(self, broker):
        """
        Get or create a connection to a broker
        """
        if (broker.host, broker.port) not in self.conns:
            self.conns[(broker.host, broker.port)] = \
                KafkaConnection(broker.host, broker.port, timeout=self.timeout)

        return self._get_conn(broker.host, broker.port)
コード例 #8
0
 def __init__(self, host, port, bufsize=4096):
     # We need one connection to bootstrap
     self.bufsize = bufsize
     self.conns = {               # (host, port) -> KafkaConnection
         (host, port): KafkaConnection(host, port, bufsize)
     }
     self.brokers = {}            # broker_id -> BrokerMetadata
     self.topics_to_brokers = {}  # topic_id -> broker_id
     self.topic_partitions = defaultdict(list)  # topic_id -> [0, 1, 2, ...]
     self._load_metadata_for_topics()
コード例 #9
0
    def test_timeout(self):
        def _timeout(*args, **kwargs):
            timeout = args[1]
            sleep(timeout)
            raise socket.timeout

        with patch.object(socket, "create_connection", side_effect=_timeout):

            with Timer() as t:
                with self.assertRaises(ConnectionError):
                    KafkaConnection("nowhere", 1234, 1.0)
            self.assertGreaterEqual(t.interval, 1.0)
コード例 #10
0
    def _get_conn(self, host, port):
        "Get or create a connection to a broker using host and port"
        host_key = (host, port)
        if self.ip_mapping is not None:
            host_key = (self._get_ip_mapping(host), port)
            log.debug('convert host address %s to %s' % (host, host_key[0]))
        if host_key not in self.conns:
            self.conns[host_key] = KafkaConnection(host,
                                                   port,
                                                   timeout=self.timeout)

        return self.conns[host_key]
コード例 #11
0
ファイル: client.py プロジェクト: mgilbir/kafka-python
 def __init__(self,
              host,
              port,
              client_id=CLIENT_ID,
              timeout=DEFAULT_SOCKET_TIMEOUT_SECONDS):
     # We need one connection to bootstrap
     self.client_id = client_id
     self.timeout = timeout
     self.conns = {  # (host, port) -> KafkaConnection
         (host, port): KafkaConnection(host, port, timeout=timeout)
     }
     self.brokers = {}  # broker_id -> BrokerMetadata
     self.topics_to_brokers = {}  # topic_id -> broker_id
     self.topic_partitions = {}  # topic_id -> [0, 1, 2, ...]
     self.load_metadata_for_topics()  # bootstrap with all metadata
コード例 #12
0
ファイル: test_conn.py プロジェクト: yobibytes/kafka-python
    def test_copy_thread(self, socket):
        """KafkaConnection copies work in other threads"""

        err = []
        copy = KafkaConnection('kafka', 9092).copy()

        def thread_func(err, copy):
            try:
                self.assertEqual(copy.host, 'kafka')
                self.assertEqual(copy.port, 9092)
                self.assertNotEqual(copy._sock, None)
            except Exception as e:
                err.append(e)
            else:
                err.append(None)
        thread = Thread(target=thread_func, args=(err, copy))
        thread.start()
        thread.join()

        self.assertEqual(err, [None])
        self.assertEqual(socket.call_count, 2)
コード例 #13
0
ファイル: test_conn.py プロジェクト: yobibytes/kafka-python
 def test_init_creates_socket_connection(self):
     KafkaConnection(self.config['host'], self.config['port'])
     self.MockCreateConn.assert_called_with((self.config['host'], self.config['port']), DEFAULT_SOCKET_TIMEOUT_SECONDS)
コード例 #14
0
ファイル: test_conn.py プロジェクト: yobibytes/kafka-python
class ConnTest(unittest.TestCase):
    def setUp(self):

        # kafka.conn debug logging is verbose, so only enable in conn tests
        logging.getLogger('kafka.conn').setLevel(logging.DEBUG)

        self.config = {
            'host': 'localhost',
            'port': 9090,
            'request_id': 0,
            'payload': b'test data',
            'payload2': b'another packet'
        }

        # Mocking socket.create_connection will cause _sock to always be a
        # MagicMock()
        patcher = mock.patch('socket.create_connection', spec=True)
        self.MockCreateConn = patcher.start()
        self.addCleanup(patcher.stop)

        # Also mock socket.sendall() to appear successful
        self.MockCreateConn().sendall.return_value = None

        # And mock socket.recv() to return two payloads, then '', then raise
        # Note that this currently ignores the num_bytes parameter to sock.recv()
        payload_size = len(self.config['payload'])
        payload2_size = len(self.config['payload2'])
        self.MockCreateConn().recv.side_effect = [
            struct.pack('>i', payload_size),
            struct.pack('>%ds' % payload_size, self.config['payload']),
            struct.pack('>i', payload2_size),
            struct.pack('>%ds' % payload2_size, self.config['payload2']),
            b''
        ]

        # Create a connection object
        self.conn = KafkaConnection(self.config['host'], self.config['port'])

        # Reset any mock counts caused by __init__
        self.MockCreateConn.reset_mock()

    def tearDown(self):
        # Return connection logging to INFO
        logging.getLogger('kafka.conn').setLevel(logging.INFO)


    def test_collect_hosts__happy_path(self):
        hosts = "localhost:1234,localhost"
        results = collect_hosts(hosts)

        self.assertEqual(set(results), set([
            ('localhost', 1234),
            ('localhost', 9092),
        ]))

    def test_collect_hosts__string_list(self):
        hosts = [
            'localhost:1234',
            'localhost',
        ]

        results = collect_hosts(hosts)

        self.assertEqual(set(results), set([
            ('localhost', 1234),
            ('localhost', 9092),
        ]))

    def test_collect_hosts__with_spaces(self):
        hosts = "localhost:1234, localhost"
        results = collect_hosts(hosts)

        self.assertEqual(set(results), set([
            ('localhost', 1234),
            ('localhost', 9092),
        ]))

    def test_send(self):
        self.conn.send(self.config['request_id'], self.config['payload'])
        self.conn._sock.sendall.assert_called_with(self.config['payload'])

    def test_init_creates_socket_connection(self):
        KafkaConnection(self.config['host'], self.config['port'])
        self.MockCreateConn.assert_called_with((self.config['host'], self.config['port']), DEFAULT_SOCKET_TIMEOUT_SECONDS)

    def test_init_failure_raises_connection_error(self):

        def raise_error(*args):
            raise socket.error

        assert socket.create_connection is self.MockCreateConn
        socket.create_connection.side_effect=raise_error
        with self.assertRaises(ConnectionError):
            KafkaConnection(self.config['host'], self.config['port'])

    def test_send__reconnects_on_dirty_conn(self):

        # Dirty the connection
        try:
            self.conn._raise_connection_error()
        except ConnectionError:
            pass

        # Now test that sending attempts to reconnect
        self.assertEqual(self.MockCreateConn.call_count, 0)
        self.conn.send(self.config['request_id'], self.config['payload'])
        self.assertEqual(self.MockCreateConn.call_count, 1)

    def test_send__failure_sets_dirty_connection(self):

        def raise_error(*args):
            raise socket.error

        assert isinstance(self.conn._sock, mock.Mock)
        self.conn._sock.sendall.side_effect=raise_error
        try:
            self.conn.send(self.config['request_id'], self.config['payload'])
        except ConnectionError:
            self.assertIsNone(self.conn._sock)

    def test_recv(self):

        self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload'])

    def test_recv__reconnects_on_dirty_conn(self):

        # Dirty the connection
        try:
            self.conn._raise_connection_error()
        except ConnectionError:
            pass

        # Now test that recv'ing attempts to reconnect
        self.assertEqual(self.MockCreateConn.call_count, 0)
        self.conn.recv(self.config['request_id'])
        self.assertEqual(self.MockCreateConn.call_count, 1)

    def test_recv__failure_sets_dirty_connection(self):

        def raise_error(*args):
            raise socket.error

        # test that recv'ing attempts to reconnect
        assert isinstance(self.conn._sock, mock.Mock)
        self.conn._sock.recv.side_effect=raise_error
        try:
            self.conn.recv(self.config['request_id'])
        except ConnectionError:
            self.assertIsNone(self.conn._sock)

    def test_recv__doesnt_consume_extra_data_in_stream(self):

        # Here just test that each call to recv will return a single payload
        self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload'])
        self.assertEqual(self.conn.recv(self.config['request_id']), self.config['payload2'])

    def test_close__object_is_reusable(self):

        # test that sending to a closed connection
        # will re-connect and send data to the socket
        self.conn.close()
        self.conn.send(self.config['request_id'], self.config['payload'])
        self.assertEqual(self.MockCreateConn.call_count, 1)
        self.conn._sock.sendall.assert_called_with(self.config['payload'])