Added sniffing logic to Transport

test_elasticsearch/test_transport.py

@@ -17,6 +17,20 @@ class DummyConnection(Connection):
             raise self.exception
         return self.status, self.data
 
+CLUSTER_NODES = '''{
+    "ok" : true,
+    "cluster_name" : "super_cluster",
+    "nodes" : {
+        "wE_6OGBNSjGksbONNncIbg" : {
+            "name" : "Nightwind",
+            "transport_address" : "inet[/127.0.0.1:9300]",
+            "hostname" : "wind",
+            "version" : "0.20.4",
+            "http_address" : "inet[/1.1.1.1:123]"
+        }
+    }
+}'''
+
 class TestTransport(TestCase):
     def test_kwargs_passed_on_to_connections(self):
         t = Transport([{'host': 'google.com'}], port=123)
@@ -47,7 +61,7 @@ class TestTransport(TestCase):
         t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection)
 
         self.assertRaises(TransportError, t.perform_request, 'GET', '/')
-        self.assertEquals(3, len(t.connection_pool.get_connection().calls))
+        self.assertEquals(3, len(t.get_connection().calls))
 
     def test_failed_connection_will_be_marked_as_dead(self):
         t = Transport([{'exception': TransportError('abandon ship')}], connection_class=DummyConnection)
@@ -55,3 +69,39 @@ class TestTransport(TestCase):
         self.assertRaises(TransportError, t.perform_request, 'GET', '/')
         self.assertEquals(0, len(t.connection_pool.connections))
 
+    def test_sniff_on_start_fetches_and_uses_nodes_list(self):
+        t = Transport([{'data': CLUSTER_NODES}], connection_class=DummyConnection, sniff_on_start=True)
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertEquals('http://1.1.1.1:123', t.get_connection().host)
+
+    def test_sniff_on_fail_triggers_sniffing_on_fail(self):
+        t = Transport([{'exception': TransportError('abandon ship')}, {"data": CLUSTER_NODES}],
+            connection_class=DummyConnection, sniff_on_connection_fail=True, max_retries=1, randomize_hosts=False)
+
+        self.assertRaises(TransportError, t.perform_request, 'GET', '/')
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertEquals('http://1.1.1.1:123', t.get_connection().host)
+
+    def test_sniff_after_n_requests(self):
+        t = Transport([{"data": CLUSTER_NODES}],
+            connection_class=DummyConnection, sniff_after_requests=5)
+
+        for _ in range(4):
+            t.perform_request('GET', '/')
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertIsInstance(t.get_connection(), DummyConnection)
+
+        t.perform_request('GET', '/')
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertEquals('http://1.1.1.1:123', t.get_connection().host)
+
+    def test_sniff_on_failure_shortens_sniff_after_n_requests(self):
+        t = Transport([{'exception': TransportError('abandon ship')}, {"data": CLUSTER_NODES}],
+            connection_class=DummyConnection, sniff_on_connection_fail=True, max_retries=1,
+            randomize_hosts=False, sniff_after_requests=4)
+
+        self.assertRaises(TransportError, t.perform_request, 'GET', '/')
+        self.assertEquals(1, len(t.connection_pool.connections))
+        self.assertEquals('http://1.1.1.1:123', t.get_connection().host)
+        self.assertEquals(3, t.sniff_after_requests)
+