Add unit/integration tests for shard aware

Author: fruch

tests/integration/standard/test_shard_aware.py

@@ -0,0 +1,117 @@
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest  # noqa
+
+from cassandra.cluster import Cluster
+from cassandra.policies import TokenAwarePolicy, RoundRobinPolicy
+
+from tests.integration import use_singledc
+
+
+def setup_module():
+    use_singledc()
+
+
+class TestShardAwareIntegration(unittest.TestCase):
+    @classmethod
+    def setup_class(cls):
+        cls.cluster = Cluster(protocol_version=4, load_balancing_policy=TokenAwarePolicy(RoundRobinPolicy()))
+        cls.session = cls.cluster.connect()
+
+    @classmethod
+    def teardown_class(cls):
+        cls.cluster.shutdown()
+
+    def verify_same_shard_in_tracing(self, results, shard_name):
+        trace_id = results.response_future.get_query_trace_ids()[0]
+        traces = self.session.execute("SELECT * FROM system_traces.events WHERE session_id = %s", (trace_id, ))
+        events = [event for event in traces]
+        for event in events:
+            print(event.thread, event.activity)
+        for event in traces:
+            self.assertEqual(event.thread, shard_name)
+        self.assertIn('querying locally', "\n".join([ event.activity for event in events ]))
+
+    def test_all_tracing_coming_one_shard(self):
+        """
+        Testing that shard aware driver is sending the requests to the correct shards
+
+        using the traces to validate that all the action been executed on the the same shard.
+        this test is using prepared SELECT statements for this validation
+        """
+
+        self.session.execute(
+            """
+            DROP KEYSPACE IF EXISTS preparedtests
+            """
+        )
+        self.session.execute(
+            """
+            CREATE KEYSPACE preparedtests
+            WITH replication = {'class': 'SimpleStrategy', 'replication_factor': '1'}
+            """)
+
+        self.session.execute("USE preparedtests")
+        self.session.execute(
+            """
+            CREATE TABLE cf0 (
+                a text,
+                b text,
+                c text,
+                PRIMARY KEY (a, b)
+            )
+            """)
+
+        prepared = self.session.prepare(
+            """
+            INSERT INTO cf0 (a, b, c) VALUES  (?, ?, ?)
+            """)
+
+        bound = prepared.bind(('a', 'b', 'c'))
+
+        self.session.execute(bound)
+
+        bound = prepared.bind(('e', 'f', 'g'))
+
+        self.session.execute(bound)
+
+        bound = prepared.bind(('100000', 'f', 'g'))
+
+        self.session.execute(bound)
+
+        prepared = self.session.prepare(
+            """
+            SELECT * FROM cf0 WHERE a=? AND b=?
+            """)
+
+        bound = prepared.bind(('a', 'b'))
+        results = self.session.execute(bound, trace=True)
+        self.assertEqual(results, [('a', 'b', 'c')])
+
+        self.verify_same_shard_in_tracing(results, "shard 4")
+
+        bound = prepared.bind(('100000', 'f'))
+        results = self.session.execute(bound, trace=True)
+        self.assertEqual(results, [('100000', 'f', 'g')])
+
+        self.verify_same_shard_in_tracing(results, "shard 2")
+
+        bound = prepared.bind(('e', 'f'))
+        results = self.session.execute(bound, trace=True)
+
+        self.verify_same_shard_in_tracing(results, "shard 4")

tests/unit/test_shard_aware.py

@@ -0,0 +1,44 @@
+# Copyright DataStax, Inc.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+try:
+    import unittest2 as unittest
+except ImportError:
+    import unittest  # noqa
+
+from cassandra.connection import ShardingInfo
+from cassandra.metadata import Murmur3Token
+
+class TestShardAware(unittest.TestCase):
+    def test_parsing_and_calculating_shard_id(self):
+        '''
+        Testing the parsing of the options command
+        and the calculation getting a shard id from a Murmur3 token
+        '''
+        class OptionsHolder():
+            options = {
+                'SCYLLA_SHARD': ['1'], 
+                'SCYLLA_NR_SHARDS': ['12'],
+                'SCYLLA_PARTITIONER': ['org.apache.cassandra.dht.Murmur3Partitioner'],
+                'SCYLLA_SHARDING_ALGORITHM': ['biased-token-round-robin'],
+                'SCYLLA_SHARDING_IGNORE_MSB': ['12']
+            }
+        shard_id, shard_info = ShardingInfo.parse_sharding_info(OptionsHolder())
+
+        self.assertEqual(shard_id, 1)
+        self.assertEqual(shard_info.shard_id(Murmur3Token.from_key(b"a")), 4)
+        self.assertEqual(shard_info.shard_id(Murmur3Token.from_key(b"b")), 6)
+        self.assertEqual(shard_info.shard_id(Murmur3Token.from_key(b"c")), 6)
+        self.assertEqual(shard_info.shard_id(Murmur3Token.from_key(b"e")), 4)
+        self.assertEqual(shard_info.shard_id(Murmur3Token.from_key(b"100000")), 2)