66import socket
77
88import pytest
9+ import time
910
1011from kafka .conn import BrokerConnection , ConnectionStates , collect_hosts
1112from kafka .protocol .api import RequestHeader
@@ -61,28 +62,99 @@ def test_connect_timeout(_socket, conn):
6162 # Initial connect returns EINPROGRESS
6263 # immediate inline connect returns EALREADY
6364 # second explicit connect returns EALREADY
64- # third explicit connect returns EALREADY and times out via last_attempt
65+ # third explicit connect returns EALREADY and times out via last_activity
6566 _socket .connect_ex .side_effect = [EINPROGRESS , EALREADY , EALREADY , EALREADY ]
6667 conn .connect ()
6768 assert conn .state is ConnectionStates .CONNECTING
6869 conn .connect ()
6970 assert conn .state is ConnectionStates .CONNECTING
71+ conn .last_activity = 0
7072 conn .last_attempt = 0
7173 conn .connect ()
7274 assert conn .state is ConnectionStates .DISCONNECTED
7375
76+ def test_connect_timeout_slowconn (_socket , conn , mocker ):
77+ # Same as test_connect_timeout,
78+ # but we make the connection run longer than the timeout in order to test that
79+ # BrokerConnection resets the timer whenever things happen during the connection
80+ # See https://github.com/dpkp/kafka-python/issues/2386
81+ _socket .connect_ex .side_effect = [EINPROGRESS , EISCONN ]
82+
83+ # 0.8 = we guarantee that when testing with three intervals of this we are past the timeout
84+ time_between_connect = (conn .config ['connection_timeout_ms' ]/ 1000 ) * 0.8
85+ start = time .time ()
86+
87+ # Use plaintext auth for simplicity
88+ last_activity = conn .last_activity
89+ last_attempt = conn .last_attempt
90+ conn .config ['security_protocol' ] = 'SASL_PLAINTEXT'
91+ conn .connect ()
92+ assert conn .state is ConnectionStates .CONNECTING
93+ # Ensure the last_activity counter was updated
94+ # Last_attempt should also be updated
95+ assert conn .last_activity > last_activity
96+ assert conn .last_attempt > last_attempt
97+ last_attempt = conn .last_attempt
98+ last_activity = conn .last_activity
99+
100+ # Simulate time being passed
101+ # This shouldn't be enough time to time out the connection
102+ conn ._try_authenticate = mocker .Mock (side_effect = [False , False , True ])
103+ with mock .patch ("time.time" , return_value = start + time_between_connect ):
104+ # This should trigger authentication
105+ # Note that an authentication attempt isn't actually made until now.
106+ # We simulate that authentication does not succeed at this point
107+ # This is technically incorrect, but it lets us see what happens
108+ # to the state machine when the state doesn't change for two function calls
109+ conn .connect ()
110+ assert conn .last_activity > last_activity
111+ # Last attempt is kept as a legacy variable, should not update
112+ assert conn .last_attempt == last_attempt
113+ last_activity = conn .last_activity
114+
115+ assert conn .state is ConnectionStates .AUTHENTICATING
116+
117+
118+ # This time around we should be way past timeout.
119+ # Now we care about connect() not terminating the attempt,
120+ # because connection state was progressed in the meantime.
121+ with mock .patch ("time.time" , return_value = start + time_between_connect * 2 ):
122+ # Simulate this one not succeeding as well. This is so we can ensure things don't time out
123+ conn .connect ()
124+
125+ # No state change = no activity change
126+ assert conn .last_activity == last_activity
127+ assert conn .last_attempt == last_attempt
128+
129+ # If last_activity was not reset when the state transitioned to AUTHENTICATING,
130+ # the connection state would be timed out now.
131+ assert conn .state is ConnectionStates .AUTHENTICATING
132+
133+
134+ # This time around, the connection should succeed.
135+ with mock .patch ("time.time" , return_value = start + time_between_connect * 3 ):
136+ # This should finalize the connection
137+ conn .connect ()
138+
139+ assert conn .last_activity > last_activity
140+ assert conn .last_attempt == last_attempt
141+ last_activity = conn .last_activity
142+
143+ assert conn .state is ConnectionStates .CONNECTED
144+
145+
74146
75147def test_blacked_out (conn ):
76148 with mock .patch ("time.time" , return_value = 1000 ):
77- conn .last_attempt = 0
149+ conn .last_activity = 0
78150 assert conn .blacked_out () is False
79- conn .last_attempt = 1000
151+ conn .last_activity = 1000
80152 assert conn .blacked_out () is True
81153
82154
83155def test_connection_delay (conn ):
84156 with mock .patch ("time.time" , return_value = 1000 ):
85- conn .last_attempt = 1000
157+ conn .last_activity = 1000
86158 assert conn .connection_delay () == conn .config ['reconnect_backoff_ms' ]
87159 conn .state = ConnectionStates .CONNECTING
88160 assert conn .connection_delay () == float ('inf' )
@@ -286,7 +358,7 @@ def test_lookup_on_connect():
286358 ]
287359
288360 with mock .patch ("socket.getaddrinfo" , return_value = mock_return2 ) as m :
289- conn .last_attempt = 0
361+ conn .last_activity = 0
290362 conn .connect ()
291363 m .assert_called_once_with (hostname , port , 0 , socket .SOCK_STREAM )
292364 assert conn ._sock_afi == afi2
@@ -301,11 +373,10 @@ def test_relookup_on_failure():
301373 assert conn .host == hostname
302374 mock_return1 = []
303375 with mock .patch ("socket.getaddrinfo" , return_value = mock_return1 ) as m :
304- last_attempt = conn .last_attempt
376+ last_activity = conn .last_activity
305377 conn .connect ()
306378 m .assert_called_once_with (hostname , port , 0 , socket .SOCK_STREAM )
307379 assert conn .disconnected ()
308- assert conn .last_attempt > last_attempt
309380
310381 afi2 = socket .AF_INET
311382 sockaddr2 = ('127.0.0.2' , 9092 )
@@ -314,12 +385,13 @@ def test_relookup_on_failure():
314385 ]
315386
316387 with mock .patch ("socket.getaddrinfo" , return_value = mock_return2 ) as m :
317- conn .last_attempt = 0
388+ conn .last_activity = 0
318389 conn .connect ()
319390 m .assert_called_once_with (hostname , port , 0 , socket .SOCK_STREAM )
320391 assert conn ._sock_afi == afi2
321392 assert conn ._sock_addr == sockaddr2
322393 conn .close ()
394+ assert conn .last_activity > last_activity
323395
324396
325397def test_requests_timed_out (conn ):
0 commit comments