Merge pull request #636 from dpkp/french_connection

Simplify connection logic

1 files changed, 19 insertions, 0 deletions

diff --git a/test/conftest.py b/test/conftest.py
index f3a8947..a389480 100644
--- a/test/conftest.py
+++ b/test/conftest.py
@@ -31,3 +31,22 @@ def kafka_broker(version, zookeeper, request):
         k.close()
     request.addfinalizer(fin)
     return k
+
+
+@pytest.fixture
+def conn(mocker):
+    from kafka.conn import ConnectionStates
+    from kafka.future import Future
+    from kafka.protocol.metadata import MetadataResponse
+    conn = mocker.patch('kafka.client_async.BrokerConnection')
+    conn.return_value = conn
+    conn.state = ConnectionStates.CONNECTED
+    conn.send.return_value = Future().success(
+        MetadataResponse[0](
+            [(0, 'foo', 12), (1, 'bar', 34)],  # brokers
+            []))  # topics
+    conn.blacked_out.return_value = False
+    conn.connect.side_effect = lambda: conn.state
+    conn.connecting = lambda: conn.connect() is ConnectionStates.CONNECTING
+    conn.connected = lambda: conn.connect() is ConnectionStates.CONNECTED
+    return conn