Skip to content

Commit 1a7f808

Browse files
author
Dana Powers
committed
Fix test_conn tests wrt _dirty / _sock
1 parent 767acb1 commit 1a7f808

File tree

1 file changed

+4
-9
lines changed

1 file changed

+4
-9
lines changed

test/test_conn.py

Lines changed: 4 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,8 @@ def setUp(self):
2525
self.MockCreateConn().sendall.return_value = None
2626
self.addCleanup(patcher.stop)
2727

28+
# And mock socket.recv() to return the payload
29+
self.MockCreateConn().recv.return_value = self.config['payload']
2830
self.conn = KafkaConnection(self.config['host'], self.config['port'])
2931
socket.create_connection.reset_mock()
3032

@@ -80,12 +82,10 @@ def raise_error(*args):
8082
def test_send__reconnects_on_dirty_conn(self):
8183

8284
# Dirty the connection
83-
assert self.conn._dirty is False
8485
try:
8586
self.conn._raise_connection_error()
8687
except ConnectionError:
8788
pass
88-
assert self.conn._dirty is True
8989

9090
# Now test that sending attempts to reconnect
9191
self.assertEqual(socket.create_connection.call_count, 0)
@@ -108,14 +108,12 @@ def test_send__failure_sets_dirty_connection(self):
108108
def raise_error(*args):
109109
raise socket.error
110110

111-
assert self.conn._dirty is False
112-
113111
assert isinstance(self.conn._sock, mock.Mock)
114112
self.conn._sock.sendall.side_effect=raise_error
115113
try:
116114
self.conn.send(self.config['request_id'], self.config['payload'])
117115
except ConnectionError:
118-
self.assertEquals(self.conn._dirty, True)
116+
self.assertIsNone(self.conn._sock)
119117

120118
def test_recv(self):
121119

@@ -142,11 +140,9 @@ def test_recv__reconnects_on_dirty_conn(self):
142140
self.conn._raise_connection_error()
143141
except ConnectionError:
144142
pass
145-
assert self.conn._dirty is True
146143

147144
# Now test that recv'ing attempts to reconnect
148145
self.assertEqual(socket.create_connection.call_count, 0)
149-
self.conn._sock.recv.return_value = self.config['payload']
150146
self.conn._read_bytes(len(self.config['payload']))
151147
self.assertEqual(socket.create_connection.call_count, 1)
152148

@@ -167,13 +163,12 @@ def raise_error(*args):
167163
raise socket.error
168164

169165
# test that recv'ing attempts to reconnect
170-
assert self.conn._dirty is False
171166
assert isinstance(self.conn._sock, mock.Mock)
172167
self.conn._sock.recv.side_effect=raise_error
173168
try:
174169
self.conn.recv(self.config['request_id'])
175170
except ConnectionError:
176-
self.assertEquals(self.conn._dirty, True)
171+
self.assertIsNone(self.conn._sock)
177172

178173
def test_recv__doesnt_consume_extra_data_in_stream(self):
179174
data1 = self.config['payload']

0 commit comments

Comments
 (0)