refactor: schema agreement test

Make test easier to read and more robust.

Author: soyacz

scylla/tests/ccm_integration/schema_agreement.rs

@@ -1,15 +1,20 @@
 use std::sync::Arc;
+use std::time::Duration;
 
 use crate::ccm::cluster::{Cluster, ClusterOptions};
 use crate::ccm::{run_ccm_test, CLUSTER_VERSION};
 use crate::common::utils::{setup_tracing, unique_keyspace_name};
 
 use scylla::client::execution_profile::ExecutionProfile;
+use scylla::client::session::Session;
 use scylla::cluster::{ClusterState, Node, NodeRef};
 use scylla::policies::load_balancing::{FallbackPlan, LoadBalancingPolicy, RoutingInfo};
 use scylla::query::Query;
 use tokio::sync::Mutex;
+use tracing::info;
+use uuid::Uuid;
 
+/// Creates a cluster configuration with 3 nodes for schema agreement tests.
 fn cluster_3_nodes() -> ClusterOptions {
     ClusterOptions {
         name: "schema_agreement_test".to_string(),
@@ -19,6 +24,7 @@ fn cluster_3_nodes() -> ClusterOptions {
     }
 }
 
+/// A load balancing policy that targets a single node.
 #[derive(Debug)]
 struct SingleTargetLBP {
     target: (Arc<Node>, Option<u32>),
@@ -46,145 +52,220 @@ impl LoadBalancingPolicy for SingleTargetLBP {
     }
 }
 
+/// Waits for schema agreement with a timeout and retries.
+async fn wait_for_schema_agreement(
+    session: &Session,
+    timeout: Duration,
+    retries: u32,
+) -> Result<Option<Uuid>, anyhow::Error> {
+    let retry_interval = Duration::from_millis(500);
+    let mut attempts = 0;
+
+    tokio::time::timeout(timeout, async {
+        loop {
+            match session.check_schema_agreement().await {
+                Ok(Some(agreement)) => return Ok(Some(agreement)),
+                Ok(None) => {
+                    attempts += 1;
+                    if attempts > retries {
+                        return Err(anyhow::anyhow!(
+                            "Schema agreement not reached after {} retries",
+                            retries
+                        ));
+                    }
+                    info!("Schema agreement not yet reached, retrying ({}/{})", attempts, retries);
+                    tokio::time::sleep(retry_interval).await;
+                }
+                Err(e) => return Err(anyhow::anyhow!("Failed to check schema agreement: {}", e)),
+            }
+        }
+    })
+    .await
+    .map_err(|_| anyhow::anyhow!("Schema agreement timed out after {:?}", timeout))?
+}
+
+/// Sets up a keyspace with a given replication factor.
+async fn setup_keyspace(
+    session: &Session,
+    keyspace: &str,
+    replication_factor: u32,
+) -> Result<(), anyhow::Error> {
+    let query = format!(
+        "CREATE KEYSPACE {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : {}}}",
+        keyspace, replication_factor
+    );
+    session.query_unpaged(query, &[]).await?;
+    session.use_keyspace(keyspace, true).await?;
+    Ok(())
+}
+
 #[tokio::test]
 #[cfg_attr(not(ccm_tests), ignore)]
 async fn test_schema_agreement() {
     setup_tracing();
     run_ccm_test(cluster_3_nodes, test_schema_agreement_all_nodes).await;
     run_ccm_test(cluster_3_nodes, test_schema_agreement_with_stopped_node).await;
     run_ccm_test(cluster_3_nodes, test_schema_agreement_with_paused_node).await;
+    // TODO - multidc cases
 }
 
+/// Tests schema agreement with all nodes running.
 async fn test_schema_agreement_all_nodes(cluster: Arc<Mutex<Cluster>>) {
     let cluster = cluster.lock().await;
-    let session = cluster.make_session_builder().await.build().await.unwrap();
+    let session = cluster
+        .make_session_builder()
+        .await
+        .build()
+        .await
+        .expect("Failed to create session");
 
-    // Create keyspace
     let keyspace = unique_keyspace_name();
+    setup_keyspace(&session, &keyspace, 3)
+        .await
+        .expect("Failed to setup keyspace");
+
+    info!("Creating table in test_schema_agreement_all_nodes");
     session
-        .query_unpaged(
-            format!(
-                "CREATE KEYSPACE {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 3}}",
-                keyspace
-            ),
-            &[],
-        )
-        .await
-        .unwrap();
-
-    // Use keyspace
-    session.use_keyspace(keyspace, true).await.unwrap();
-
-    // Create a table and check schema agreement
-    let _result = session
-        .query_unpaged(
-            "CREATE TABLE test_schema_agreement_all (k int primary key, v int)",
-            &[],
-        )
-        .await
-        .unwrap();
-
-    // Check if schema is in agreement
-    let schema_agreement = session.check_schema_agreement().await.unwrap();
-    assert!(schema_agreement.is_some());
+        .query_unpaged("CREATE TABLE test_table (k int primary key, v int)", &[])
+        .await
+        .expect("Failed to create table");
+
+    let agreement = wait_for_schema_agreement(&session, Duration::from_secs(10), 20)
+        .await
+        .expect("Schema agreement failed");
+    assert!(agreement.is_some(), "Schema agreement should be reached");
+    info!("Schema agreement achieved with all nodes");
 }
 
+/// Tests schema agreement with one node stopped.
 async fn test_schema_agreement_with_stopped_node(cluster: Arc<Mutex<Cluster>>) {
     let cluster = cluster.lock().await;
-
-    // Create keyspace
-    let session = cluster.make_session_builder().await.build().await.unwrap();
+    let session = cluster
+        .make_session_builder()
+        .await
+        .build()
+        .await
+        .expect("Failed to create session");
 
     let keyspace = unique_keyspace_name();
+    setup_keyspace(&session, &keyspace, 3)
+        .await
+        .expect("Failed to setup keyspace");
+
+    let node = cluster
+        .nodes()
+        .get_by_id(2)
+        .await
+        .expect("Failed to get node 2");
+    info!("Stopping node 2");
+    node.write().await.stop(None).await.expect("Failed to stop node");
+
+    info!("Creating table with one node stopped");
     session
-        .query_unpaged(
-            format!(
-                "CREATE KEYSPACE {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 3}}",
-                keyspace
-            ),
-            &[],
-        )
-        .await
-        .unwrap();
-
-    // Use keyspace
-    session.use_keyspace(keyspace, true).await.unwrap();
-
-    // Stop node 2
-    let node = cluster.nodes().get_by_id(2).await.unwrap();
-    node.write().await.stop(None).await.unwrap();
-
-    // Create a table while one node is stopped
-    let _result = session
-        .query_unpaged(
-            "CREATE TABLE test_schema_agreement_stopped (k int primary key, v int)",
-            &[],
-        )
-        .await
-        .unwrap();
-
-    // Schema agreement should succeed with remaining up nodes
-    let schema_agreement = session.check_schema_agreement().await.unwrap();
-    assert!(schema_agreement.is_some());
-
-    // Start the node back
-    node.write().await.start(None).await.unwrap();
-    let schema_agreement = session.check_schema_agreement().await.unwrap();
-    assert!(schema_agreement.is_some());
+        .query_unpaged("CREATE TABLE test_table (k int primary key, v int)", &[])
+        .await
+        .expect("Failed to create table");
+
+    let agreement = wait_for_schema_agreement(&session, Duration::from_secs(10), 20)
+        .await
+        .expect("Schema agreement failed with stopped node");
+    assert!(
+        agreement.is_some(),
+        "Schema agreement should be reached with remaining nodes"
+    );
+
+    info!("Restarting node 2");
+    node.write()
+        .await
+        .start(None)
+        .await
+        .expect("Failed to restart node");
+    let agreement = wait_for_schema_agreement(&session, Duration::from_secs(10), 20)
+        .await
+        .expect("Schema agreement failed after restart");
+    assert!(
+        agreement.is_some(),
+        "Schema agreement should be reached after node restart"
+    );
+    info!("Schema agreement achieved after node restart");
 }
 
+/// Tests schema agreement with one node paused.
 async fn test_schema_agreement_with_paused_node(cluster: Arc<Mutex<Cluster>>) {
     let cluster = cluster.lock().await;
-
-    let session = cluster.make_session_builder().await.build().await.unwrap();
+    let session = cluster
+        .make_session_builder()
+        .await
+        .build()
+        .await
+        .expect("Failed to create session");
 
     let keyspace = unique_keyspace_name();
-    session
-        .query_unpaged(
-            format!(
-                "CREATE KEYSPACE {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 3}}",
-                keyspace
-            ),
-            &[],
-        )
+    setup_keyspace(&session, &keyspace, 3)
         .await
-        .unwrap();
-
-    session.use_keyspace(keyspace, true).await.unwrap();
+        .expect("Failed to setup keyspace");
 
-    // Stop node 2
     let node_id = 2;
-    let ccm_node = cluster.nodes().get_by_id(node_id).await.unwrap();
+    let ccm_node = cluster
+        .nodes()
+        .get_by_id(node_id)
+        .await
+        .expect("Failed to get node 2");
     let ccm_node_addr = ccm_node.read().await.broadcast_rpc_address().clone();
-    ccm_node.write().await.pause().await.unwrap();
+    info!("Pausing node 2");
+    ccm_node
+        .write()
+        .await
+        .pause()
+        .await
+        .expect("Failed to pause node");
 
-    // Find the corresponding Scylla node from the session to avoid querying it directly
     let cluster_state = session.get_cluster_state();
-    let scylla_node = cluster_state
+    let running_scylla_node = cluster_state
         .get_nodes_info()
         .iter()
         .find(|n| n.address.ip() != ccm_node_addr)
-        .expect("Could not find unpaused Scylla node for querying");
+        .expect("Could not find unpaused Scylla node");
 
     let policy = SingleTargetLBP {
-        target: (scylla_node.clone(), Some(0)),
+        target: (running_scylla_node.clone(), Some(0)),
     };
     let execution_profile = ExecutionProfile::builder()
         .load_balancing_policy(Arc::new(policy))
         .build();
-    let mut stmt =
-        Query::new("CREATE TABLE test_schema_agreement_paused (k int primary key, v int)");
+    let mut stmt = Query::new("CREATE TABLE test_table (k int primary key, v int)");
     stmt.set_execution_profile_handle(Some(execution_profile.into_handle()));
-    // Create a table while one node is paused
-    let _result = session.query_unpaged(stmt, &[]).await.unwrap();
 
-    // Schema agreement should succeed with remaining up nodes
-    let schema_agreement = session.check_schema_agreement().await.unwrap();
-    assert!(schema_agreement.is_some());
+    info!("Creating table with one node paused");
+    session
+        .query_unpaged(stmt, &[])
+        .await
+        .expect("Failed to create table");
+
+    let agreement = wait_for_schema_agreement(&session, Duration::from_secs(10), 20)
+        .await
+        .expect("Schema agreement failed with paused node");
+    assert!(
+        agreement.is_some(),
+        "Schema agreement should be reached with remaining nodes"
+    );
 
-    // Start the node back
-    ccm_node.write().await.resume().await.unwrap();
+    info!("Resuming node 2");
+    ccm_node
+        .write()
+        .await
+        .resume()
+        .await
+        .expect("Failed to resume node");
 
-    let schema_agreement = session.check_schema_agreement().await.unwrap();
-    assert!(schema_agreement.is_some());
+    let agreement = wait_for_schema_agreement(&session, Duration::from_secs(10), 20)
+        .await
+        .expect("Schema agreement failed after resume");
+    assert!(
+        agreement.is_some(),
+        "Schema agreement should be reached after node resume"
+    );
+    info!("Schema agreement achieved after node resume");
 }
+
+