Refactor dns lookup in BrokerConnection (#1312)

1 files changed, 25 insertions, 0 deletions

diff --git a/test/test_conn.py b/test/test_conn.py
index 1621e60..ef7925a 100644
--- a/test/test_conn.py
+++ b/test/test_conn.py
@@ -267,3 +267,28 @@ def test_lookup_on_connect():
         m.assert_called_once_with(hostname, port, 0, 1)
         conn.close()
         assert conn.host == ip2
+
+
+def test_relookup_on_failure():
+    hostname = 'example.org'
+    port = 9092
+    conn = BrokerConnection(hostname, port, socket.AF_UNSPEC)
+    assert conn.host == conn.hostname == hostname
+    mock_return1 = []
+    with mock.patch("socket.getaddrinfo", return_value=mock_return1) as m:
+        last_attempt = conn.last_attempt
+        conn.connect()
+        m.assert_called_once_with(hostname, port, 0, 1)
+        assert conn.disconnected()
+        assert conn.last_attempt > last_attempt
+
+    ip2 = '127.0.0.2'
+    mock_return2 = [
+        (2, 2, 17, '', (ip2, 9092)),
+    ]
+
+    with mock.patch("socket.getaddrinfo", return_value=mock_return2) as m:
+        conn.connect()
+        m.assert_called_once_with(hostname, port, 0, 1)
+        conn.close()
+        assert conn.host == ip2