1

Author: dkropachev

tests/unit/test_shard_aware.py

@@ -138,12 +138,19 @@ def mock_connection_factory(self, *args, **kwargs):
                         assert connection.endpoint == DefaultEndPoint("1.2.3.4", port=port)
 
                 sleep_time = 0
-                schedule = reconnection_policy.new_schedule()
-                for _ in range(shard_count+1):
-                    sleep_time += next(schedule)
-
-                if sleep_time > 0:
-                    time.sleep(sleep_time)
+                if reconnection_policy:
+                    # Check that connections to shards are being established according to the policy
+                    # Calculate total time it will need to establish all connections
+                    # Sleep half of the time and check that connections are not there yet
+                    # Sleep rest of the time + 1 second and check that all connections has been established
+                    schedule = reconnection_policy.new_schedule()
+                    for _ in range(shard_count):
+                        sleep_time += next(schedule)
+                    if sleep_time > 0:
+                        time.sleep(sleep_time/2)
+                        # Check that connection are not being established quicker than expected
+                        assert len(pool._connections) < expected_after
+                        time.sleep(sleep_time/2 + 1)
 
                 assert len(pool._connections) == expected_after
         finally: