Move callback processing from BrokerConnection to KafkaClientpending_completions

1 files changed, 2 insertions, 1 deletions

diff --git a/test/test_client.py b/test/test_client.py
index 42a1623..d02c621 100644
--- a/test/test_client.py
+++ b/test/test_client.py
@@ -28,6 +28,7 @@ def mock_conn(conn, success=True):
     else:
         mocked.send.return_value = Future().failure(Exception())
     conn.return_value = mocked
+    conn.recv.return_value = []
 
 
 class TestSimpleClient(unittest.TestCase):
@@ -94,7 +95,7 @@ class TestSimpleClient(unittest.TestCase):
         mock_conn(mocked_conns[('kafka03', 9092)], success=False)
         future = Future()
         mocked_conns[('kafka02', 9092)].send.return_value = future
-        mocked_conns[('kafka02', 9092)].recv.side_effect = lambda: future.success('valid response')
+        mocked_conns[('kafka02', 9092)].recv.return_value = [('valid response', future)]
 
         def mock_get_conn(host, port, afi):
             return mocked_conns[(host, port)]