caching_session: Test the new batch methods

Add a test which tests the newly added methods
prepare_batch and batch in CachingSession.

Signed-off-by: Jan Ciolek <[email protected]>

scylla/src/transport/caching_session.rs

@@ -148,8 +148,13 @@ impl CachingSession {
 #[cfg(test)]
 mod tests {
     use crate::utils::test_utils::unique_keyspace_name;
-    use crate::{CachingSession, Session, SessionBuilder};
+    use crate::{
+        batch::{Batch, BatchStatement},
+        prepared_statement::PreparedStatement,
+        CachingSession, Session, SessionBuilder,
+    };
     use futures::TryStreamExt;
+    use std::collections::BTreeSet;
 
     async fn new_for_test() -> Session {
         let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
@@ -290,4 +295,144 @@ mod tests {
         assert_eq!(1, session.cache.len());
         assert_eq!(1, result.rows.unwrap().len());
     }
+
+    async fn assert_test_batch_table_rows_contain(
+        sess: &CachingSession,
+        expected_rows: &[(i32, i32)],
+    ) {
+        let selected_rows: BTreeSet<(i32, i32)> = sess
+            .execute("SELECT a, b FROM test_batch_table", ())
+            .await
+            .unwrap()
+            .rows_typed::<(i32, i32)>()
+            .unwrap()
+            .map(|r| r.unwrap())
+            .collect();
+        for expected_row in expected_rows.iter() {
+            if !selected_rows.contains(expected_row) {
+                panic!(
+                    "Expected {:?} to contain row: {:?}, but they didnt",
+                    selected_rows, expected_row
+                );
+            }
+        }
+    }
+
+    #[tokio::test]
+    async fn test_batch() {
+        let session: CachingSession = create_caching_session().await;
+
+        session
+            .execute(
+                "CREATE TABLE IF NOT EXISTS test_batch_table (a int, b int, primary key (a, b))",
+                (),
+            )
+            .await
+            .unwrap();
+
+        let unprepared_insert_a_b: &str = "insert into test_batch_table (a, b) values (?, ?)";
+        let unprepared_insert_a_7: &str = "insert into test_batch_table (a, b) values (?, 7)";
+        let unprepared_insert_8_b: &str = "insert into test_batch_table (a, b) values (8, ?)";
+        let prepared_insert_a_b: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_a_b.into())
+            .await
+            .unwrap();
+        let prepared_insert_a_7: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_a_7.into())
+            .await
+            .unwrap();
+        let prepared_insert_8_b: PreparedStatement = session
+            .add_prepared_statement(&unprepared_insert_8_b.into())
+            .await
+            .unwrap();
+
+        let assert_batch_prepared = |b: &Batch| {
+            for stmt in &b.statements {
+                match stmt {
+                    BatchStatement::PreparedStatement(_) => {}
+                    _ => panic!("Unprepared statement in prepared batch!"),
+                }
+            }
+        };
+
+        {
+            let mut unprepared_batch: Batch = Default::default();
+            unprepared_batch.append_statement(unprepared_insert_a_b);
+            unprepared_batch.append_statement(unprepared_insert_a_7);
+            unprepared_batch.append_statement(unprepared_insert_8_b);
+
+            session
+                .batch(&unprepared_batch, ((10, 20), (10,), (20,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(10, 20), (10, 7), (8, 20)]).await;
+
+            let prepared_batch: Batch = session.prepare_batch(&unprepared_batch).await.unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((15, 25), (15,), (25,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(15, 25), (15, 7), (8, 25)]).await;
+        }
+
+        {
+            let mut partially_prepared_batch: Batch = Default::default();
+            partially_prepared_batch.append_statement(unprepared_insert_a_b);
+            partially_prepared_batch.append_statement(prepared_insert_a_7.clone());
+            partially_prepared_batch.append_statement(unprepared_insert_8_b);
+
+            session
+                .batch(&partially_prepared_batch, ((30, 40), (30,), (40,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(30, 40), (30, 7), (8, 40)]).await;
+
+            let prepared_batch: Batch = session
+                .prepare_batch(&partially_prepared_batch)
+                .await
+                .unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((35, 45), (35,), (45,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(35, 45), (35, 7), (8, 45)]).await;
+        }
+
+        {
+            let mut fully_prepared_batch: Batch = Default::default();
+            fully_prepared_batch.append_statement(prepared_insert_a_b);
+            fully_prepared_batch.append_statement(prepared_insert_a_7);
+            fully_prepared_batch.append_statement(prepared_insert_8_b);
+
+            session
+                .batch(&fully_prepared_batch, ((50, 60), (50,), (60,)))
+                .await
+                .unwrap();
+            assert_test_batch_table_rows_contain(&session, &[(50, 60), (50, 7), (8, 60)]).await;
+
+            let prepared_batch: Batch = session.prepare_batch(&fully_prepared_batch).await.unwrap();
+            assert_batch_prepared(&prepared_batch);
+
+            session
+                .batch(&prepared_batch, ((55, 65), (55,), (65,)))
+                .await
+                .unwrap();
+
+            assert_test_batch_table_rows_contain(&session, &[(55, 65), (55, 7), (8, 65)]).await;
+        }
+
+        {
+            let mut bad_batch: Batch = Default::default();
+            bad_batch.append_statement(unprepared_insert_a_b);
+            bad_batch.append_statement("This isnt even CQL");
+            bad_batch.append_statement(unprepared_insert_8_b);
+
+            assert!(session.batch(&bad_batch, ((1, 2), (), (2,))).await.is_err());
+            assert!(session.prepare_batch(&bad_batch).await.is_err());
+        }
+    }
 }