Fix bug in socket timeout per PR #161 by maciejkula, add test

2 files changed, 14 insertions, 4 deletions

diff --git a/kafka/conn.py b/kafka/conn.py
index 4fdeb17..5dc3d5a 100644
--- a/kafka/conn.py
+++ b/kafka/conn.py
@@ -150,6 +150,6 @@ class KafkaConnection(local):
         """
         self.close()
         self._sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-        self._sock.connect((self.host, self.port))
         self._sock.settimeout(self.timeout)
+        self._sock.connect((self.host, self.port))
         self._dirty = False
diff --git a/test/test_client_integration.py b/test/test_client_integration.py
index db6cac9..e566cce 100644
--- a/test/test_client_integration.py
+++ b/test/test_client_integration.py
@@ -1,9 +1,10 @@
 import unittest
 import time
+import socket
+import random
 
-from kafka import *  # noqa
-from kafka.common import *  # noqa
-from kafka.codec import has_gzip, has_snappy
+import kafka
+from kafka.common import *
 from .fixtures import ZookeeperFixture, KafkaFixture
 from .testutil import *
 
@@ -19,6 +20,15 @@ class TestKafkaClientIntegration(KafkaIntegrationTestCase):
         cls.server.close()
         cls.zk.close()
 
+    def test_timeout(self):
+        server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+        server_socket.bind(('localhost', 14567))
+
+        with Timer() as t:
+            with self.assertRaises((socket.timeout, socket.error)):
+                conn = kafka.conn.KafkaConnection("localhost", 14567, 1.0)
+        self.assertGreaterEqual(t.interval, 1.0)
+
     def test_consume_none(self):
         fetch = FetchRequest(self.topic, 0, 0, 1024)