Implementa integration tests for shard connection backof policies

Tests cover:
1. LimitedConcurrencyShardConnectionBackoffPolicy
2. NoDelayShardConnectionBackoffPolicy

For both Scylla and Cassandra backend.

Author: dkropachev

tests/integration/long/test_policies.py

@@ -11,17 +11,30 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
+import os
+import time
 import unittest
+from typing import Optional
+from unittest.mock import Mock
 
 from cassandra import ConsistencyLevel, Unavailable
-from cassandra.cluster import ExecutionProfile, EXEC_PROFILE_DEFAULT
+from cassandra.cluster import ExecutionProfile, EXEC_PROFILE_DEFAULT, Session
+from cassandra.policies import LimitedConcurrencyShardConnectionBackoffPolicy, ConstantReconnectionPolicy, \
+    ShardConnectionBackoffPolicy, NoDelayShardConnectionBackoffPolicy, _ScopeBucket, \
+    _NoDelayShardConnectionBackoffScheduler
+from cassandra.shard_info import _ShardingInfo
 
-from tests.integration import use_cluster, get_cluster, get_node, TestCluster
+from tests.integration import use_cluster, get_cluster, get_node, TestCluster, remove_cluster
 
 
 def setup_module():
-    use_cluster('test_cluster', [4])
+    os.environ['SCYLLA_EXT_OPTS'] = "--smp 8"
+    use_cluster('test_cluster', [2])
+
+
+def teardown_module():
+    remove_cluster()
+    del os.environ['SCYLLA_EXT_OPTS']
 
 
 class RetryPolicyTests(unittest.TestCase):
@@ -65,3 +78,90 @@ def test_should_rethrow_on_unvailable_with_default_policy_if_cas(self):
         self.assertEqual(exception.consistency, ConsistencyLevel.SERIAL)
         self.assertEqual(exception.required_replicas, 2)
         self.assertEqual(exception.alive_replicas, 1)
+
+
+class ShardBackoffPolicyTests(unittest.TestCase):
+    def test_limited_concurrency_1_connection_per_host(self):
+        self._test_backoff(
+            LimitedConcurrencyShardConnectionBackoffPolicy(
+                backoff_policy=ConstantReconnectionPolicy(0.1),
+                max_concurrent=1,
+            )
+        )
+
+    def test_limited_concurrency_2_connection_per_host(self):
+        self._test_backoff(
+            LimitedConcurrencyShardConnectionBackoffPolicy(
+                backoff_policy=ConstantReconnectionPolicy(0.1),
+                max_concurrent=1,
+            )
+        )
+
+    def test_no_delay(self):
+        self._test_backoff(NoDelayShardConnectionBackoffPolicy())
+
+    def _test_backoff(self, shard_connection_backoff_policy: ShardConnectionBackoffPolicy):
+        backoff_policy = None
+        if isinstance(shard_connection_backoff_policy, LimitedConcurrencyShardConnectionBackoffPolicy):
+            backoff_policy = shard_connection_backoff_policy.backoff_policy
+
+        cluster = TestCluster(
+            shard_connection_backoff_policy=shard_connection_backoff_policy,
+            reconnection_policy=ConstantReconnectionPolicy(0),
+        )
+
+        # Collect scheduled calls and execute them right away
+        scheduler_calls = []
+        original_schedule = cluster.scheduler.schedule
+        pending = 0
+
+        def new_schedule(delay, fn, *args, **kwargs):
+            nonlocal pending
+            pending+=1
+
+            def execute():
+                nonlocal pending
+                try:
+                    fn(*args, **kwargs)
+                finally:
+                    pending-=1
+
+            scheduler_calls.append((delay, fn, args, kwargs))
+            return original_schedule(0, execute)
+
+        cluster.scheduler.schedule = Mock(side_effect=new_schedule)
+
+        session = cluster.connect()
+        sharding_info = get_sharding_info(session)
+
+        # Since scheduled calls executed in a separate thread we need to give them some time to complete
+        while pending > 0:
+            time.sleep(0.01)
+
+        if not sharding_info:
+            # If it is not scylla `ShardConnectionBackoffScheduler` should not be involved
+            for delay, fn, args, kwargs in scheduler_calls:
+                if fn.__self__.__class__ is _ScopeBucket or fn.__self__.__class__ is _NoDelayShardConnectionBackoffScheduler:
+                    self.fail(
+                        "in non-shard-aware case connection should be created directly, not involving ShardConnectionBackoffScheduler")
+            return
+
+        sleep_time = 0
+        if backoff_policy:
+            schedule = backoff_policy.new_schedule()
+            sleep_time = next(iter(schedule))
+
+        # Make sure that all scheduled calls have delay according to policy
+        found_related_calls = 0
+        for delay, fn, args, kwargs in scheduler_calls:
+            if fn.__self__.__class__ is _ScopeBucket or fn.__self__.__class__ is _NoDelayShardConnectionBackoffScheduler:
+                found_related_calls += 1
+                self.assertEqual(delay, sleep_time)
+        self.assertLessEqual(len(session.hosts) * (sharding_info.shards_count - 1), found_related_calls)
+
+
+def get_sharding_info(session: Session) -> Optional[_ShardingInfo]:
+    for host in session.hosts:
+        if host.sharding_info:
+            return host.sharding_info
+    return None