PYTHON-871 Added test simulating network isolation

beltran · beltran · commit d56fede690bb · 2018-01-30T16:24:14.000-05:00
diff --git a/tests/integration/simulacron/test_connection.py b/tests/integration/simulacron/test_connection.py
@@ -36,7 +36,8 @@
                                                 prime_query, prime_request,
                                                 start_and_prime_cluster_defaults,
                                                 start_and_prime_singledc,
-                                                clear_queries)
+                                                clear_queries, RejectConnections,
+                                                RejectType, AcceptConnections)
 
 
 class TrackDownListener(HostStateListener):
@@ -46,6 +47,15 @@ def __init__(self):
     def on_down(self, host):
         self.hosts_marked_down.append(host)
 
+    def on_up(self, host):
+        pass
+
+    def on_add(self, host):
+        pass
+
+    def on_remove(self, host):
+        pass
+
 class ThreadTracker(ThreadPoolExecutor):
     called_functions = []
 
@@ -339,3 +349,37 @@ class ExtendedConnection(AsyncoreConnection):
                           connection_class=ExtendedConnection)
         cluster.connect()
         cluster.shutdown()
+
+    def test_driver_recovers_nework_isolation(self):
+        start_and_prime_singledc()
+
+        idle_heartbeat_timeout = 3
+        idle_heartbeat_interval = 1
+
+        listener = TrackDownListener()
+
+        cluster = Cluster(['127.0.0.1'],
+                          load_balancing_policy=RoundRobinPolicy(),
+                          idle_heartbeat_timeout=idle_heartbeat_timeout,
+                          idle_heartbeat_interval=idle_heartbeat_interval,
+                          executor_threads=16)
+        session = cluster.connect(wait_for_all_pools=True)
+
+        cluster.register_listener(listener)
+
+        prime_request(PrimeOptions(then=NO_THEN))
+        prime_request(RejectConnections(RejectType.REJECT_STARTUP))
+
+        time.sleep((idle_heartbeat_timeout + idle_heartbeat_interval) * 2)
+
+        for host in cluster.metadata.all_hosts():
+            self.assertIn(host, listener.hosts_marked_down)
+
+        self.assertRaises(NoHostAvailable, session.execute, "SELECT * from system.local")
+
+        clear_queries()
+        prime_request(AcceptConnections())
+
+        time.sleep(idle_heartbeat_timeout + idle_heartbeat_interval + 2)
+
+        self.assertIsNotNone(session.execute("SELECT * from system.local"))
diff --git a/tests/integration/simulacron/utils.py b/tests/integration/simulacron/utils.py
@@ -103,7 +103,7 @@ def submit_request(self, query):
 
         request = Request("http://{}/{}{}".format(
             self.admin_addr, query.path, query.fetch_url_params()), data=data)
-        request.get_method = lambda: 'POST'
+        request.get_method = lambda: query.method
         request.add_header("Content-Type", 'application/json')
         request.add_header("Content-Length", len(data))
 
@@ -142,9 +142,13 @@ def clear_all_queries(self, cluster_name=DEFAULT_CLUSTER):
 
 class SimulacronRequest(object):
     def fetch_json(self):
-        raise NotImplementedError()
+        return {}
 
     def fetch_url_params(self):
+        return ""
+
+    @property
+    def method(self):
         raise NotImplementedError()
 
 
@@ -175,6 +179,44 @@ def fetch_json(self):
     def fetch_url_params(self):
         return ""
 
+    @property
+    def method(self):
+        return "POST"
+
+
+class RejectType():
+    UNBIND = "UNBIND"
+    STOP = "STOP"
+    REJECT_STARTUP = "REJECT_STARTUP"
+
+
+class RejectConnections(SimulacronRequest):
+    """
+    Class used for making simulacron reject new connections
+    """
+    def __init__(self, reject_type, cluster_name=DEFAULT_CLUSTER):
+        self.path = "listener/{}".format(cluster_name)
+        self.reject_type = reject_type
+
+    def fetch_url_params(self):
+        return "?type={0}".format(self.reject_type)
+
+    @property
+    def method(self):
+        return "DELETE"
+
+
+class AcceptConnections(SimulacronRequest):
+    """
+    Class used for making simulacron reject new connections
+    """
+    def __init__(self, cluster_name=DEFAULT_CLUSTER):
+        self.path = "listener/{}".format(cluster_name)
+
+    @property
+    def method(self):
+        return "PUT"
+
 
 class PrimeQuery(SimulacronRequest):
     """
@@ -228,6 +270,9 @@ def set_node(self, cluster_id, datacenter_id, node_id):
     def fetch_url_params(self):
         return ""
 
+    @property
+    def method(self):
+        return "POST"
 
 class ClusterQuery(SimulacronRequest):
     """
@@ -251,6 +296,9 @@ def fetch_url_params(self):
         return "?cassandra_version={0}&data_centers={1}&name={2}".\
             format(self.cassandra_version, self.data_centers, self.cluster_name)
 
+    @property
+    def method(self):
+        return "POST"
 
 def prime_driver_defaults():
     """
