@@ -25,6 +25,8 @@ def setUp(self):
25
25
self .MockCreateConn ().sendall .return_value = None
26
26
self .addCleanup (patcher .stop )
27
27
28
+ # And mock socket.recv() to return the payload
29
+ self .MockCreateConn ().recv .return_value = self .config ['payload' ]
28
30
self .conn = KafkaConnection (self .config ['host' ], self .config ['port' ])
29
31
socket .create_connection .reset_mock ()
30
32
@@ -80,12 +82,10 @@ def raise_error(*args):
80
82
def test_send__reconnects_on_dirty_conn (self ):
81
83
82
84
# Dirty the connection
83
- assert self .conn ._dirty is False
84
85
try :
85
86
self .conn ._raise_connection_error ()
86
87
except ConnectionError :
87
88
pass
88
- assert self .conn ._dirty is True
89
89
90
90
# Now test that sending attempts to reconnect
91
91
self .assertEqual (socket .create_connection .call_count , 0 )
@@ -108,14 +108,12 @@ def test_send__failure_sets_dirty_connection(self):
108
108
def raise_error (* args ):
109
109
raise socket .error
110
110
111
- assert self .conn ._dirty is False
112
-
113
111
assert isinstance (self .conn ._sock , mock .Mock )
114
112
self .conn ._sock .sendall .side_effect = raise_error
115
113
try :
116
114
self .conn .send (self .config ['request_id' ], self .config ['payload' ])
117
115
except ConnectionError :
118
- self .assertEquals (self .conn ._dirty , True )
116
+ self .assertIsNone (self .conn ._sock )
119
117
120
118
def test_recv (self ):
121
119
@@ -142,11 +140,9 @@ def test_recv__reconnects_on_dirty_conn(self):
142
140
self .conn ._raise_connection_error ()
143
141
except ConnectionError :
144
142
pass
145
- assert self .conn ._dirty is True
146
143
147
144
# Now test that recv'ing attempts to reconnect
148
145
self .assertEqual (socket .create_connection .call_count , 0 )
149
- self .conn ._sock .recv .return_value = self .config ['payload' ]
150
146
self .conn ._read_bytes (len (self .config ['payload' ]))
151
147
self .assertEqual (socket .create_connection .call_count , 1 )
152
148
@@ -167,13 +163,12 @@ def raise_error(*args):
167
163
raise socket .error
168
164
169
165
# test that recv'ing attempts to reconnect
170
- assert self .conn ._dirty is False
171
166
assert isinstance (self .conn ._sock , mock .Mock )
172
167
self .conn ._sock .recv .side_effect = raise_error
173
168
try :
174
169
self .conn .recv (self .config ['request_id' ])
175
170
except ConnectionError :
176
- self .assertEquals (self .conn ._dirty , True )
171
+ self .assertIsNone (self .conn ._sock )
177
172
178
173
def test_recv__doesnt_consume_extra_data_in_stream (self ):
179
174
data1 = self .config ['payload' ]
0 commit comments