Merge pull request #1351 from wprzytula/reorganise-integration-tests

Reorganise integration tests - part 2

Author: wprzytula

scylla/tests/integration/load_balancing/latency_awareness.rs

@@ -0,0 +1,26 @@
+use scylla::client::execution_profile::ExecutionProfile;
+use scylla::policies::load_balancing::{DefaultPolicy, LatencyAwarenessBuilder};
+
+use crate::utils::{create_new_session_builder, setup_tracing};
+
+// This is a regression test for #696.
+#[tokio::test]
+#[ntest::timeout(1000)]
+async fn latency_aware_query_completes() {
+    setup_tracing();
+    let policy = DefaultPolicy::builder()
+        .latency_awareness(LatencyAwarenessBuilder::default())
+        .build();
+    let handle = ExecutionProfile::builder()
+        .load_balancing_policy(policy)
+        .build()
+        .into_handle();
+
+    let session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
+    session.query_unpaged("whatever", ()).await.unwrap_err();
+}

scylla/tests/integration/load_balancing/load_balancing.rs

@@ -1,5 +1,6 @@
-#[expect(clippy::module_inception)]
-mod load_balancing;
+mod latency_awareness;
 mod lwt_optimisation;
 mod shards;
+mod simple_strategy;
 mod tablets;
+mod token_awareness;

scylla/tests/integration/load_balancing/mod.rs

@@ -2,10 +2,14 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use crate::utils::{
-    scylla_supports_tablets, setup_tracing, test_with_3_node_cluster, unique_keyspace_name,
-    PerformDDL,
+    create_new_session_builder, scylla_supports_tablets, setup_tracing, test_with_3_node_cluster,
+    unique_keyspace_name, PerformDDL,
 };
+use scylla::client::execution_profile::ExecutionProfile;
 use scylla::client::session_builder::SessionBuilder;
+use scylla::cluster::{ClusterState, NodeRef};
+use scylla::policies::load_balancing::{FallbackPlan, LoadBalancingPolicy, RoutingInfo};
+use scylla::routing::Shard;
 use tokio::sync::mpsc;
 
 use scylla_proxy::TargetShard;
@@ -87,3 +91,54 @@ async fn test_consistent_shard_awareness() {
         Err(err) => panic!("{}", err),
     }
 }
+
+/// Regression test for panic that appeared if LBP picked a
+/// shard that is out of range for the node.
+#[tokio::test]
+async fn test_shard_out_of_range() {
+    setup_tracing();
+
+    #[derive(Debug)]
+    struct ShardOutOfRangeLBP;
+    impl LoadBalancingPolicy for ShardOutOfRangeLBP {
+        fn pick<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            cluster: &'a ClusterState,
+        ) -> Option<(NodeRef<'a>, Option<Shard>)> {
+            let node = &cluster.get_nodes_info()[0];
+            match node.sharder() {
+                Some(sharder) => Some((node, Some(sharder.nr_shards.get() as u32))),
+                // For Cassandra let's pick some crazy shard number - it should be ignored anyway.
+                None => Some((node, Some(u16::MAX as u32))),
+            }
+        }
+
+        fn fallback<'a>(
+            &'a self,
+            _query: &'a RoutingInfo,
+            _cluster: &'a ClusterState,
+        ) -> FallbackPlan<'a> {
+            Box::new(std::iter::empty())
+        }
+
+        fn name(&self) -> String {
+            "ShardOutOfRangeLBP".into()
+        }
+    }
+
+    let handle = ExecutionProfile::builder()
+        .load_balancing_policy(Arc::new(ShardOutOfRangeLBP))
+        .build()
+        .into_handle();
+    let session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
+    let _ = session
+        .query_unpaged("SELECT * FROM system.local WHERE key='local'", ())
+        .await
+        .unwrap();
+}

scylla/tests/integration/load_balancing/shards.rs

@@ -0,0 +1,64 @@
+use crate::utils::{create_new_session_builder, unique_keyspace_name, PerformDDL as _};
+
+/// It's recommended to use NetworkTopologyStrategy everywhere, so most tests use only NetworkTopologyStrategy.
+/// We still support SimpleStrategy, so to make sure that SimpleStrategy works correctly this test runs
+/// a few queries in a SimpleStrategy keyspace.
+#[tokio::test]
+async fn simple_strategy_test() {
+    let ks = unique_keyspace_name();
+    let session = create_new_session_builder().build().await.unwrap();
+
+    session
+        .ddl(format!(
+            "CREATE KEYSPACE {} WITH REPLICATION = \
+                {{'class': 'SimpleStrategy', 'replication_factor': 1}}",
+            ks
+        ))
+        .await
+        .unwrap();
+
+    session
+        .ddl(format!(
+            "CREATE TABLE {}.tab (p int, c int, r int, PRIMARY KEY (p, c, r))",
+            ks
+        ))
+        .await
+        .unwrap();
+
+    session
+        .query_unpaged(
+            format!("INSERT INTO {}.tab (p, c, r) VALUES (1, 2, 3)", ks),
+            (),
+        )
+        .await
+        .unwrap();
+
+    session
+        .query_unpaged(
+            format!("INSERT INTO {}.tab (p, c, r) VALUES (?, ?, ?)", ks),
+            (4, 5, 6),
+        )
+        .await
+        .unwrap();
+
+    let prepared = session
+        .prepare(format!("INSERT INTO {}.tab (p, c, r) VALUES (?, ?, ?)", ks))
+        .await
+        .unwrap();
+
+    session.execute_unpaged(&prepared, (7, 8, 9)).await.unwrap();
+
+    let mut rows: Vec<(i32, i32, i32)> = session
+        .query_unpaged(format!("SELECT p, c, r FROM {}.tab", ks), ())
+        .await
+        .unwrap()
+        .into_rows_result()
+        .unwrap()
+        .rows::<(i32, i32, i32)>()
+        .unwrap()
+        .map(|r| r.unwrap())
+        .collect::<Vec<(i32, i32, i32)>>();
+    rows.sort();
+
+    assert_eq!(rows, vec![(1, 2, 3), (4, 5, 6), (7, 8, 9)]);
+}

scylla/tests/integration/load_balancing/simple_strategy.rs

@@ -0,0 +1,66 @@
+use crate::utils::{
+    create_new_session_builder, scylla_supports_tablets, setup_tracing, unique_keyspace_name,
+    PerformDDL as _,
+};
+
+#[tokio::test]
+async fn test_token_awareness() {
+    setup_tracing();
+    let session = create_new_session_builder().build().await.unwrap();
+    let ks = unique_keyspace_name();
+
+    // Need to disable tablets in this test because they make token routing
+    // work differently, and in this test we want to test the classic token ring
+    // behavior.
+    let mut create_ks = format!(
+        "CREATE KEYSPACE IF NOT EXISTS {ks} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}"
+    );
+    if scylla_supports_tablets(&session).await {
+        create_ks += " AND TABLETS = {'enabled': false}"
+    }
+
+    session.ddl(create_ks).await.unwrap();
+    session
+        .ddl(format!(
+            "CREATE TABLE IF NOT EXISTS {}.t (a text primary key)",
+            ks
+        ))
+        .await
+        .unwrap();
+
+    let mut prepared_statement = session
+        .prepare(format!("INSERT INTO {}.t (a) VALUES (?)", ks))
+        .await
+        .unwrap();
+    prepared_statement.set_tracing(true);
+
+    // The default policy should be token aware
+    for size in 1..50usize {
+        let key = vec!['a'; size].into_iter().collect::<String>();
+        let values = (&key,);
+
+        // Execute a query and observe tracing info
+        let res = session
+            .execute_unpaged(&prepared_statement, values)
+            .await
+            .unwrap();
+        let tracing_info = session
+            .get_tracing_info(res.tracing_id().as_ref().unwrap())
+            .await
+            .unwrap();
+
+        // Verify that only one node was involved
+        assert_eq!(tracing_info.nodes().len(), 1);
+
+        // Do the same with execute_iter (it now works with writes)
+        let iter = session
+            .execute_iter(prepared_statement.clone(), values)
+            .await
+            .unwrap();
+        let tracing_id = iter.tracing_ids()[0];
+        let tracing_info = session.get_tracing_info(&tracing_id).await.unwrap();
+
+        // Again, verify that only one node was involved
+        assert_eq!(tracing_info.nodes().len(), 1);
+    }
+}

scylla/tests/integration/load_balancing/token_awareness.rs

@@ -0,0 +1,55 @@
+use crate::utils::{
+    create_new_session_builder, setup_tracing, unique_keyspace_name, PerformDDL as _,
+};
+
+#[tokio::test]
+async fn test_macros_complex_pk() {
+    setup_tracing();
+    let session = create_new_session_builder().build().await.unwrap();
+    let ks = unique_keyspace_name();
+
+    session.ddl(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks)).await.unwrap();
+    session.use_keyspace(ks, true).await.unwrap();
+    session
+        .ddl("CREATE TABLE IF NOT EXISTS complex_pk (a int, b int, c text, d int, e int, primary key ((a,b,c),d))")
+        .await
+        .unwrap();
+
+    // Check that SerializeRow and DeserializeRow macros work
+    {
+        #[derive(scylla::SerializeRow, scylla::DeserializeRow, PartialEq, Debug, Clone)]
+        struct ComplexPk {
+            a: i32,
+            b: i32,
+            c: Option<String>,
+            d: i32,
+            e: i32,
+        }
+        let input: ComplexPk = ComplexPk {
+            a: 9,
+            b: 8,
+            c: Some("seven".into()),
+            d: 6,
+            e: 5,
+        };
+        session
+            .query_unpaged(
+                "INSERT INTO complex_pk (a,b,c,d,e) VALUES (?,?,?,?,?)",
+                input.clone(),
+            )
+            .await
+            .unwrap();
+        let output: ComplexPk = session
+            .query_unpaged(
+                "SELECT * FROM complex_pk WHERE a = 9 and b = 8 and c = 'seven'",
+                &[],
+            )
+            .await
+            .unwrap()
+            .into_rows_result()
+            .unwrap()
+            .single_row()
+            .unwrap();
+        assert_eq!(input, output)
+    }
+}

scylla/tests/integration/macros/complex_pk.rs

@@ -1 +1,2 @@
+mod complex_pk;
 mod hygiene;