diff --git a/kafka/net/asyncio_backend.py b/kafka/net/asyncio_backend.py index 24f7628ff..4cbbf9cdf 100644 --- a/kafka/net/asyncio_backend.py +++ b/kafka/net/asyncio_backend.py @@ -329,99 +329,64 @@ async def create_connection(self, protocol, host, port, *, ssl=None, 'The asyncio backend does not support proxy_url yet; use the ' 'default selector backend for SOCKS5/HTTP-CONNECT proxying.') server_hostname = host.rstrip('.') if ssl is not None else None - adapter = _AsyncioProtocolAdapter() + adapter = _AsyncioProtocolAdapter(protocol, host, port, socket_options) connect = self._loop.create_connection( lambda: adapter, host, port, ssl=ssl, server_hostname=server_hostname) try: if timeout_at is not None: connect = asyncio.wait_for(connect, max(0.0, timeout_at - time.monotonic())) - aio_transport, _ = await connect + await connect + if adapter.error is not None: + raise adapter.error except asyncio.TimeoutError: raise Errors.KafkaConnectionError('Connection timed out') except Errors.KafkaError: raise except Exception as exc: # noqa: BLE001 -- surface any connect error uniformly raise Errors.KafkaConnectionError('unable to connect to %s:%s: %s' % (host, port, exc)) + + +class _AsyncioProtocolAdapter(asyncio.Protocol): + """Thin asyncio.Protocol that wires a KafkaConnection to a wrapped transport.""" + + def __init__(self, conn, host, port, socket_options=()): + self._conn = conn + self._host = host + self._port = port + self._socket_options = socket_options + self.error = None + self.transport = None # the _AsyncioTransport wrapper + + def connection_made(self, aio_transport): sock = aio_transport.get_extra_info('socket') if sock is not None: - for option in socket_options: + for option in self._socket_options: try: sock.setsockopt(*option) except OSError: pass - return _AsyncioTransport(aio_transport, adapter, host, port) - - -class _AsyncioProtocolAdapter(asyncio.Protocol): - """Bridges asyncio's Protocol callbacks to a KafkaConnection. - - asyncio calls ``connection_made`` on this adapter during - ``create_connection`` -- before the manager wires the KafkaConnection via - ``transport.set_protocol(conn)`` (Option A: the caller runs connection_made - after its "closed during connect" check). Reading is paused until wired, so - no data is delivered early; a buffer/latched-loss is kept as a safety net. - """ - - def __init__(self): - self._kafka = None # the KafkaConnection, once wired - self._transport = None - self._buffer = [] - self._eof = False - self._lost = False - self._lost_exc = None - self._paused_writing = False - self._on_read = None # bumps the wrapper's last_read - - def connection_made(self, transport): - self._transport = transport - # Hold off delivery until the KafkaConnection is wired + resumes reading. - transport.pause_reading() + self.transport = _AsyncioTransport(aio_transport, self._host, self._port) + try: + self._conn.connection_made(self.transport) + except Exception as exc: # noqa: BLE001 -- conn refused (closed mid-connect) + self.error = exc + aio_transport.abort() def data_received(self, data): - if self._on_read is not None: - self._on_read() - if self._kafka is None: - self._buffer.append(data) - else: - self._kafka.data_received(data) + self.transport._bump_read() + self._conn.data_received(data) def eof_received(self): - if self._kafka is not None: - return self._kafka.eof_received() - self._eof = True - return None + return self._conn.eof_received() def connection_lost(self, exc): - if self._kafka is not None: - self._kafka.connection_lost(exc) - else: - self._lost = True - self._lost_exc = exc + self._conn.connection_lost(exc) def pause_writing(self): - if self._kafka is not None: - self._kafka.pause_writing() - else: - self._paused_writing = True + self._conn.pause_writing() def resume_writing(self): - if self._kafka is not None: - self._kafka.resume_writing() - else: - self._paused_writing = False - - def _wire(self, kafka_protocol, on_read): - self._kafka = kafka_protocol - self._on_read = on_read - if self._paused_writing: - kafka_protocol.pause_writing() - for data in self._buffer: - kafka_protocol.data_received(data) - self._buffer = [] - if self._eof: - kafka_protocol.eof_received() - if self._lost: - kafka_protocol.connection_lost(self._lost_exc) + self._conn.resume_writing() class _AsyncioTransport: @@ -433,9 +398,8 @@ class _AsyncioTransport: protocol, host/host_port/getPeer, and last_activity for idle sweeping. """ - def __init__(self, transport, adapter, host, port): + def __init__(self, transport, host, port): self._t = transport - self._adapter = adapter self.host = host self._port = port self._protocol = None @@ -454,11 +418,11 @@ def get_protocol(self): def set_protocol(self, protocol): self._protocol = protocol - self._adapter._wire(protocol, self._bump_read) def write(self, data): self.last_write = time.monotonic() self._t.write(data) + return len(data) def close(self): self._t.close() diff --git a/kafka/net/backend.py b/kafka/net/backend.py index 0bbb94de3..4b7ad7bfb 100644 --- a/kafka/net/backend.py +++ b/kafka/net/backend.py @@ -108,10 +108,15 @@ class Transport(Protocol): """The transport surface a backend's ``create_connection`` returns. The subset of the ``asyncio.Transport`` / Twisted ``ITransport`` surface - that ``KafkaConnection`` actually drives. The selector's - ``KafkaTCPTransport`` and an asyncio-transport adapter both satisfy it. + that ``KafkaConnection`` actually drives (plus ``last_activity``, which the + manager reads to sweep idle connections). A backend's ``create_connection`` + builds one and wires it to the conn. The selector's ``KafkaTCPTransport`` + and an asyncio-transport adapter both satisfy it. """ + # Monotonic timestamp of the last read/write; manager idle-sweeping reads it. + last_activity: float + def write(self, data: bytes) -> None: ... def close(self) -> None: ... def abort(self, error: Any = None) -> None: ... @@ -177,17 +182,20 @@ async def create_connection( proxy_url: Optional[str] = None, socket_options: Sequence[Any] = (), timeout_at: Optional[float] = None, - ) -> Transport: - """Establish and return a connected :class:`Transport` to ``host:port``. - - The backend owns DNS, connect, TLS and (where supported) proxying. The - *caller* wires the transport to the protocol afterwards via - ``protocol.connection_made(transport)`` (so manager-level policy such as - the "closed during connect" check runs first). ``protocol`` (a - ``KafkaConnection``) is passed because backends that own the socket - (asyncio/Twisted) need it at connect time to receive transport events; - they buffer inbound data until ``connection_made`` is called. Backends - without native proxy support raise when ``proxy_url`` is set. + ) -> None: + """Establish a connected :class:`Transport` to ``host:port`` and wire it. + + The backend owns DNS, connect, TLS and (where supported) proxying, builds + a :class:`Transport`, and wires it to ``protocol`` by calling + ``protocol.connection_made(transport)`` itself -- mirroring + ``asyncio.loop.create_connection`` / Twisted, which own the socket and + wire the protocol at connect time. Nothing is returned: the caller drives + the connection through ``protocol`` (``conn.transport``), never a + transport handle. ``protocol`` (a ``KafkaConnection``) may *refuse* the + transport by raising from ``connection_made`` if it closed mid-connect; + on that (or any) failure the backend closes the orphaned transport + before propagating. Backends without native proxy support raise when + ``proxy_url`` is set. """ # --- cross-thread bridge --------------------------------------------- diff --git a/kafka/net/connection.py b/kafka/net/connection.py index 45ca27f8d..1107ce164 100644 --- a/kafka/net/connection.py +++ b/kafka/net/connection.py @@ -168,7 +168,9 @@ def _send_request(self, request, future=None, timeout_at=None): # in_flight_requests (len==1), trip the >= check, pause, and never be # written to the transport - hanging forever. if not self.paused: - self.transport.write(self.parser.send_bytes()) + total_bytes = self.transport.write(self.parser.send_bytes()) + if self._sensors: + self._sensors.bytes_sent.record(total_bytes) if len(self.in_flight_requests) >= self.config['max_in_flight_requests_per_connection']: self.pause('max_in_flight') return future @@ -193,6 +195,8 @@ def data_received(self, data): if self.closed: log.debug('%s: Ignoring %d bytes received by closed connection', self, len(data)) return + if self._sensors: + self._sensors.bytes_received.record(len(data)) responses = self.parser.receive_bytes(data) # augment responses w/ correlation_id, future, and timestamp @@ -301,7 +305,9 @@ def unpause(self, v): if not self.paused and self.parser and self.transport: to_send = self.parser.send_bytes() if to_send: - self.transport.write(to_send) + total_bytes = self.transport.write(to_send) + if self._sensors: + self._sensors.bytes_sent.record(total_bytes) def pause_writing(self): """ Called when the transport's buffer goes over the high-water mark. diff --git a/kafka/net/manager.py b/kafka/net/manager.py index ef538e451..68091ab26 100644 --- a/kafka/net/manager.py +++ b/kafka/net/manager.py @@ -236,31 +236,24 @@ def ssl_enabled(self): return self.config['security_protocol'] in ('SSL', 'SASL_SSL') async def _connect(self, node, conn, reset_backoff_on_connect=True, timeout_at=None): - # Tracks ownership of the freshly built transport: while non-None it is - # ours to clean up (the connection hasn't taken it over yet), so the - # finally clause closes it. Cleared once connection_made() succeeds. - transport = None try: - transport = await self._net.create_connection( + await self._net.create_connection( conn, node.host, node.port, ssl=self.ssl_context, proxy_url=self.config['proxy_url'], socket_options=self.config['socket_options'], timeout_at=timeout_at) - # The connection (or the whole manager) may have been closed while - # we were building the transport. Handing it to connection_made() - # would flip the conn back to `initializing` and resurrect a - # connection that is already being torn down. Discard - # the new transport instead of reviving a dead connection. - if conn.closed or self.closed: - log.debug('%s: closed during connect; discarding new transport', conn) - return - conn.connection_made(transport) - transport = None # conn owns cleanup now; skip finally: transport.close() # Note: conn.initialize does not currently raise on error; # errors are pushed to conn.init_future and raised on await conn await conn.initialize(timeout_at=timeout_at) except Exception as exc: + if conn.closed or self.closed: + # A concurrent close() raced the connect (manager / bootstrap + # teardown). connection_made() refused to resurrect the conn and + # the backend already discarded the transport; don't back off a + # connection that is going away. + log.debug('%s: closed during connect; discarding', conn) + return log.error('Connection failed: %s', exc) conn.connection_lost(exc) self.update_backoff(node.node_id) @@ -268,9 +261,6 @@ async def _connect(self, node, conn, reset_backoff_on_connect=True, timeout_at=N Errors.AuthorizationError)): self._auth_failures[node.node_id] = exc return - finally: - if transport is not None: - transport.close() if self._sensors: self._sensors.connection_created.record() diff --git a/kafka/net/selector.py b/kafka/net/selector.py index 7f11daa4f..3ec6dc30e 100644 --- a/kafka/net/selector.py +++ b/kafka/net/selector.py @@ -490,15 +490,16 @@ def create_future(self): async def create_connection(self, protocol, host, port, *, ssl=None, proxy_url=None, socket_options=(), timeout_at=None): - """Establish and return a connected transport to host:port. + """Establish a connected transport to host:port and wire ``protocol``. The selector owns the raw socket: DNS + non-blocking connect (with optional SOCKS5/HTTP-CONNECT proxy via KafkaNetSocket), then wraps it - in a TCP or SSL transport and runs the TLS handshake. ``protocol`` (the - KafkaConnection) is not used here -- the caller wires it via - ``connection_made()`` after its own "closed during connect" check; it - is part of the contract because socket-owning backends (asyncio, - Twisted) need it at connect time. + in a TCP or SSL transport, runs the TLS handshake, and calls + ``protocol.connection_made(transport)`` -- mirroring asyncio/Twisted, + which own the socket and wire the protocol at connect time. On any + failure (handshake error, or a ``protocol`` that refuses the transport + because it closed mid-connect) the transport is closed before raising, + so the caller never handles a transport instance directly. """ sock = await _inet_create_connection(self, host, port, socket_options, proxy_url=proxy_url, timeout_at=timeout_at) @@ -509,8 +510,13 @@ async def create_connection(self, protocol, host, port, *, ssl=None, try: await transport.handshake() except Exception as e: + transport.close() raise Errors.KafkaConnectionError('Handshake failed: %s' % e) - return transport + try: + protocol.connection_made(transport) + except Exception: + transport.close() + raise def sleep(self, delay): return KernelEvent('_sleep', delay) diff --git a/kafka/net/transport.py b/kafka/net/transport.py index f44d0ff20..339d8198f 100644 --- a/kafka/net/transport.py +++ b/kafka/net/transport.py @@ -95,8 +95,6 @@ async def _read_from_sock(self): return self.abort(error=err) log.debug('%s: received %d bytes', self, len(recvd_data)) self.last_read = time.monotonic() - if self._protocol and self._protocol._sensors: - self._protocol._sensors.bytes_received.record(len(recvd_data)) try: self._protocol.data_received(recvd_data) except Errors.KafkaProtocolError as e: @@ -177,6 +175,7 @@ def write(self, data): if not self._writing: self._writing = True self._write_task = self._net.call_soon(self._write_to_sock) + return len(data) def writelines(self, list_of_data): """Write a list (or any iterable) of data bytes to the transport.""" @@ -186,6 +185,7 @@ def writelines(self, list_of_data): if not self._writing: self._writing = True self._write_task = self._net.call_soon(self._write_to_sock) + return sum(len(data) for data in list_of_data) async def _write_to_sock(self): try: @@ -196,8 +196,6 @@ async def _write_to_sock(self): return self.abort(error=err) log.debug('%s: sent %d bytes', self, total_bytes) self.last_write = time.monotonic() - if self._protocol and self._protocol._sensors: - self._protocol._sensors.bytes_sent.record(total_bytes) finally: self._writing = False if self._closed: diff --git a/test/conftest.py b/test/conftest.py index a64b30003..a0040543d 100644 --- a/test/conftest.py +++ b/test/conftest.py @@ -58,18 +58,22 @@ def client(net, manager, broker): @pytest.fixture def net(): - return NetworkSelector() + backend = NetworkSelector() + try: + yield backend + finally: + backend.close() @pytest.fixture def manager(net, broker): + broker.attach(net) manager = KafkaConnectionManager( net, bootstrap_servers='%s:%d' % (broker.host, broker.port), api_version=broker.broker_version, request_timeout_ms=5000, ) - broker.attach(manager) try: yield manager finally: diff --git a/test/consumer/test_coordinator.py b/test/consumer/test_coordinator.py index bb6a2c4c8..78489896d 100644 --- a/test/consumer/test_coordinator.py +++ b/test/consumer/test_coordinator.py @@ -45,6 +45,10 @@ def coordinator(broker, client, metrics): try: yield coord finally: + # Drop any group generation a test left set so close() doesn't fire a + # LeaveGroupRequest the MockBroker isn't scripted for (which otherwise + # surfaces as teardown ERROR-log noise + an unhandled mock-broker task). + coord.reset_generation() coord.close(timeout_ms=0) @@ -468,7 +472,7 @@ async def fake_send_retry_then_success(_partitions): assert coordinator._send_offset_fetch_request.call_count == 4 # successful, failed, retried+success -def test_close(mocker, coordinator): +def test_close(mocker, net, coordinator): mocker.patch.object(coordinator, '_maybe_auto_commit_offsets_sync') mocker.patch.object(coordinator, '_handle_leave_group_response') mocker.patch.object(coordinator, 'coordinator_unknown', return_value=False) @@ -476,7 +480,8 @@ def test_close(mocker, coordinator): coordinator._generation = Generation(1, 'foobar', b'') coordinator.state = MemberState.STABLE cli = coordinator._client - mocker.patch.object(cli._manager, 'send', return_value=Future().success('foobar')) + mocker.patch.object(cli._manager, 'send', + return_value=net.create_future().success('foobar')) mocker.patch.object(cli, 'poll') coordinator.close() @@ -1352,7 +1357,7 @@ def test_join_group_async_raises_non_retriable(request, broker, seeded_coord): def test_join_group_async_returns_false_on_short_timeout_and_caches_task( - request, broker, seeded_coord): + request, broker, net, seeded_coord): """Short consumer.poll(timeout_ms=N) should return False instead of hanging when the broker is slow to respond to JoinGroup; the in-flight task is cached so the next poll re-awaits it instead of sending a fresh @@ -1370,7 +1375,7 @@ def test_join_group_async_returns_false_on_short_timeout_and_caches_task( seeded_coord.state = MemberState.UNJOINED # JoinGroup response future controlled by the test. Hangs until released. - join_response_pending = seeded_coord._manager.create_future() # awaited by slow_join_handler + join_response_pending = net.create_future() # awaited by slow_join_handler join_request_count = [0] async def slow_join_handler(api_key, api_version, correlation_id, request_bytes): @@ -1456,10 +1461,9 @@ def test_ensure_active_group_sync_facade(request, broker, seeded_coord): assert seeded_coord.state == MemberState.STABLE -def test_heartbeat(mocker, coordinator): +def test_heartbeat(mocker, net, coordinator): coordinator.coordinator_id = 0 coordinator.state = MemberState.STABLE - net = coordinator._manager._net assert not coordinator._heartbeat_enabled and not coordinator._heartbeat_closed @@ -1489,7 +1493,7 @@ def test_heartbeat(mocker, coordinator): # spin); using side_effect with an async function that awaits the Future # forces the suspension we want. The Mock's call_count then verifies the # loop fired exactly once. - blocked_send = coordinator._manager.create_future() + blocked_send = net.create_future() async def _hang(*args, **kwargs): await blocked_send mocker.patch.object(coordinator, '_send_heartbeat_request', side_effect=_hang) @@ -2336,8 +2340,9 @@ def _stub_leave_group(self, mocker, coordinator): return_value=False) coordinator.coordinator_id = 0 cli = coordinator._client + net = cli._net mocker.patch.object(cli._manager, 'send', - return_value=Future().success('foobar')) + return_value=net.create_future().success('foobar')) mocker.patch.object(cli, 'poll') def test_close_revokes_for_live_group(self, mocker, coordinator): diff --git a/test/mock_broker.py b/test/mock_broker.py index 7c5b7f3d2..3061ad113 100644 --- a/test/mock_broker.py +++ b/test/mock_broker.py @@ -13,8 +13,8 @@ broker.respond(JoinGroupRequest, JoinGroupResponse(version=5, ...)) broker.respond(SyncGroupRequest, SyncGroupResponse(version=3, ...)) - # Attach to a KafkaConnectionManager so new connections use MockTransport: - broker.attach(manager) + # Attach to a NetBackend so new connections use MockTransport: + broker.attach(net) # Or use the consumer/client factory fixtures for a one-liner setup. @@ -128,6 +128,7 @@ def write(self, data): self._write_buffer.extend(data) self.last_write = time.monotonic() self._net.call_soon(self._process_requests) + return len(data) def writelines(self, data_list): for data in data_list: @@ -488,16 +489,16 @@ async def _resolve_response(response, api_key, api_version, correlation_id, requ response = await response return response - def attach(self, manager): - """Monkey-patch a KafkaConnectionManager to route all connections + def attach(self, net): + """Monkey-patch a NetBackend to route all connections through this MockBroker. - After calling this, any ``manager.get_connection(node_id)`` call will - create a KafkaConnection backed by a MockTransport connected to this + After calling this, any ``net.create_connection`` call will + connect a KafkaConnection to a MockTransport linked with this broker, instead of opening a real TCP socket. Arguments: - manager: A ``KafkaConnectionManager`` instance. + net: A ``NetBackend`` instance. """ broker = self @@ -506,11 +507,20 @@ async def _mock_create_connection(protocol, host, port, **kwargs): raise Errors.KafkaConnectionError( 'connect to %s:%s refused (MockBroker stopped)' % (host, port)) - return MockTransport( - manager._net, broker, + transport = MockTransport( + net, broker, node_id=protocol.node_id, host=host, port=port) + # create_connection contract: the backend wires the protocol and + # cleans up the transport if the conn refuses (closed mid-connect). + # Nothing is returned -- the conn drives it via conn.transport. + try: + protocol.connection_made(transport) + except Exception: + transport.close() + raise - manager._net.create_connection = _mock_create_connection + net.create_connection = _mock_create_connection + return net def client_factory(self): """Return a callable suitable for passing as ``kafka_client=...`` @@ -531,7 +541,7 @@ def client_factory(self): def factory(**kwargs): from kafka.net.compat import KafkaNetClient client = KafkaNetClient(**kwargs) - broker.attach(client._manager) + broker.attach(client._net) return client return factory @@ -704,9 +714,17 @@ async def _mock_create_connection(protocol, host, port, **kwargs): if broker is None or not broker.online: raise Errors.KafkaConnectionError( 'connect to %s:%s refused' % (host, port)) - return MockTransport( + transport = MockTransport( manager._net, broker, node_id=protocol.node_id, host=host, port=port) + # create_connection contract: the backend wires the protocol and + # cleans up the transport if the conn refuses (closed mid-connect). + # Nothing is returned -- the conn drives it via conn.transport. + try: + protocol.connection_made(transport) + except Exception: + transport.close() + raise manager._net.create_connection = _mock_create_connection diff --git a/test/net/test_asyncio_backend.py b/test/net/test_asyncio_backend.py index 06c8c2971..aea1ebdae 100644 --- a/test/net/test_asyncio_backend.py +++ b/test/net/test_asyncio_backend.py @@ -12,6 +12,7 @@ import pytest +import kafka.errors as Errors from kafka.net.asyncio_backend import AsyncioBackend, AsyncioFuture from kafka.net.backend import NetBackend from kafka.net.manager import KafkaConnectionManager @@ -182,13 +183,20 @@ class _StubProtocol: """Minimal KafkaConnection-shaped protocol for transport tests.""" def __init__(self): self.received = bytearray() - self.lost = False + self.closed = False + self.transport = None + def connection_made(self, transport): + # Mirror KafkaConnection.connection_made: adopt + wire + start reading. + # create_connection() calls this itself under the new backend contract. + self.transport = transport + transport.set_protocol(self) + transport.resume_reading() def data_received(self, data): self.received += data def eof_received(self): return None def connection_lost(self, exc): - self.lost = True + self.closed = True def pause_writing(self): pass def resume_writing(self): @@ -225,10 +233,10 @@ def test_roundtrip_write_read_close(self, started_backend): proto = _StubProtocol() async def do(): - transport = await started_backend.create_connection(proto, host, port) - # Wire like manager._connect -> conn.connection_made does. - transport.set_protocol(proto) - transport.resume_reading() + # create_connection wires proto via proto.connection_made() itself + # and returns nothing; the transport lives on the conn (proto). + await started_backend.create_connection(proto, host, port) + transport = proto.transport assert transport.host_port() == '%s:%s' % (host, port) assert transport.getPeer()[0:2] == (host, port) transport.write(b'ping') @@ -245,8 +253,9 @@ async def do(): srv.close() def test_early_data_is_buffered_until_wired(self, started_backend): - # Server sends immediately on connect; data must not be lost before - # the protocol is wired via set_protocol(). + # Server sends immediately on connect; data must not be lost. The adapter + # buffers between asyncio's connection_made and the wiring done inside + # create_connection (proto.connection_made -> set_protocol flushes it). srv = socket.socket(socket.AF_INET, socket.SOCK_STREAM) srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) srv.bind(('127.0.0.1', 0)) @@ -263,10 +272,9 @@ def serve(): proto = _StubProtocol() async def do(): - transport = await started_backend.create_connection(proto, host, port) - await started_backend.sleep(0.05) # let early bytes arrive (buffered) - transport.set_protocol(proto) # wiring flushes the buffer - transport.resume_reading() + # create_connection wires proto (flushing any early-buffered bytes). + await started_backend.create_connection(proto, host, port) + transport = proto.transport for _ in range(100): if proto.received: break @@ -289,6 +297,30 @@ async def do(): with pytest.raises(Exception): # KafkaConnectionError started_backend.run(do) + def test_connection_made_refusal_is_raised(self, started_backend): + # A conn that closed mid-connect refuses the transport by raising from + # connection_made. create_connection must surface that exception (not + # swallow it) and not leak the transport. + srv, host, port = _echo_server() + + class _RefusingProtocol(_StubProtocol): + def connection_made(self, transport): + self.transport = transport + raise Errors.KafkaConnectionError('Connection closed during connect') + + proto = _RefusingProtocol() + + async def do(): + await started_backend.create_connection(proto, host, port) + + try: + with pytest.raises(Errors.KafkaConnectionError, match='closed during connect'): + started_backend.run(do) + # The transport was aborted, so it is closing (not leaked). + assert proto.transport.is_closing() + finally: + srv.close() + def test_proxy_url_raises(self, started_backend): async def do(): await started_backend.create_connection( @@ -329,13 +361,13 @@ class TestEndToEndMockBroker: def test_bootstrap_and_send(self): broker = MockBroker() net = AsyncioBackend(client_id='e2e') + broker.attach(net) manager = KafkaConnectionManager( net, bootstrap_servers='%s:%d' % (broker.host, broker.port), api_version=broker.broker_version, request_timeout_ms=5000, ) - broker.attach(manager) net.start() try: manager.bootstrap(timeout_ms=5000) diff --git a/test/net/test_connection.py b/test/net/test_connection.py index 005f9ae3b..b458c3be0 100644 --- a/test/net/test_connection.py +++ b/test/net/test_connection.py @@ -471,7 +471,7 @@ def test_sasl_authenticate_handshake_error(self, net): handshake_response.error_code = 33 # UnsupportedSaslMechanismError handshake_response.mechanisms = ['GSSAPI'] - f = Future() + f = net.create_future() f.success(handshake_response) conn._send_request = MagicMock(return_value=f) @@ -496,7 +496,7 @@ def test_sasl_authenticate_mechanism_not_supported(self, net): handshake_response.error_code = 0 handshake_response.mechanisms = ['GSSAPI', 'SCRAM-SHA-256'] - f = Future() + f = net.create_future() f.success(handshake_response) conn._send_request = MagicMock(return_value=f) diff --git a/test/net/test_manager.py b/test/net/test_manager.py index ebf60eac4..99450391c 100644 --- a/test/net/test_manager.py +++ b/test/net/test_manager.py @@ -438,9 +438,15 @@ def test_connect_discards_transport_when_closed_during_build(self, net): transport = MagicMock() async def fake_create_connection(protocol, host, port, **kwargs): - # Simulate a concurrent close landing mid-connect. + # Simulate a concurrent close landing mid-connect, then the new + # create_connection contract: the backend wires the protocol (which + # refuses, since the conn closed) and cleans up the transport. conn.close() - return transport + try: + protocol.connection_made(transport) + except Exception: + transport.close() + raise with patch.object(net, 'create_connection', side_effect=fake_create_connection): diff --git a/test/net/test_sasl_reauthentication.py b/test/net/test_sasl_reauthentication.py index fe6ba4612..76745c23a 100644 --- a/test/net/test_sasl_reauthentication.py +++ b/test/net/test_sasl_reauthentication.py @@ -41,6 +41,7 @@ def sasl_broker(): @pytest.fixture def sasl_manager(net, sasl_broker): + sasl_broker.attach(net) manager = KafkaConnectionManager( net, bootstrap_servers='%s:%d' % (sasl_broker.host, sasl_broker.port), @@ -48,7 +49,6 @@ def sasl_manager(net, sasl_broker): request_timeout_ms=5000, **SASL_CONFIG, ) - sasl_broker.attach(manager) try: yield manager finally: diff --git a/test/producer/test_transaction_manager_mock_broker.py b/test/producer/test_transaction_manager_mock_broker.py index d040507d4..f9699d3c2 100644 --- a/test/producer/test_transaction_manager_mock_broker.py +++ b/test/producer/test_transaction_manager_mock_broker.py @@ -64,7 +64,7 @@ def _make_client(broker): request_timeout_ms=5000, metadata_max_age_ms=300000, ) - broker.attach(client._manager) + broker.attach(client._net) # Bootstrap so cluster metadata has the broker node and we have an # api_version mapping available for subsequent sends. client.check_version(timeout_ms=5000) diff --git a/test/test_cluster.py b/test/test_cluster.py index 50b1486e0..e4a4b726d 100644 --- a/test/test_cluster.py +++ b/test/test_cluster.py @@ -499,7 +499,8 @@ def test_request_update_sends_metadata_request(self, manager, net): manager.cluster.config['retry_backoff_ms'] = 10 # reduce loop delay when metadata in progress response = _make_metadata_response(8) - with patch.object(manager, 'send', return_value=Future().success(response)): + with patch.object(manager, 'send', + return_value=net.create_future().success(response)): f = manager.cluster.request_update() # Drive the cluster refresh loop net.poll(timeout_ms=100, future=f) diff --git a/test/test_mock_broker.py b/test/test_mock_broker.py index 6c2c1df76..0f20f9eca 100644 --- a/test/test_mock_broker.py +++ b/test/test_mock_broker.py @@ -128,7 +128,7 @@ def _make_client(self, broker): request_timeout_ms=5000, metadata_max_age_ms=300000, ) - broker.attach(client._manager) + broker.attach(client._net) return client def test_bootstrap_through_mock(self): @@ -303,7 +303,7 @@ def test_api_version_negotiation(self): request_timeout_ms=5000, metadata_max_age_ms=300000, ) - broker.attach(client._manager) + broker.attach(client._net) try: version = client.check_version(timeout_ms=5000) assert version == broker.broker_version