From fad88a135110f3e79d20499d039a570740d74b10 Mon Sep 17 00:00:00 2001 From: sylwiaszunejko Date: Tue, 26 Sep 2023 11:54:44 +0200 Subject: [PATCH] transport: add test for preparing batch --- scylla/src/transport/mod.rs | 2 + .../transport/silent_prepare_batch_test.rs | 110 ++++++++++++++++++ 2 files changed, 112 insertions(+) create mode 100644 scylla/src/transport/silent_prepare_batch_test.rs diff --git a/scylla/src/transport/mod.rs b/scylla/src/transport/mod.rs index 6025bf639f..939983cfc4 100644 --- a/scylla/src/transport/mod.rs +++ b/scylla/src/transport/mod.rs @@ -28,6 +28,8 @@ mod authenticate_test; mod cql_collections_test; #[cfg(test)] mod session_test; +#[cfg(test)] +mod silent_prepare_batch_test; #[cfg(test)] mod cql_types_test; diff --git a/scylla/src/transport/silent_prepare_batch_test.rs b/scylla/src/transport/silent_prepare_batch_test.rs new file mode 100644 index 0000000000..70286369de --- /dev/null +++ b/scylla/src/transport/silent_prepare_batch_test.rs @@ -0,0 +1,110 @@ +use crate::{ + batch::{Batch, BatchStatement}, + prepared_statement::PreparedStatement, + test_utils::{create_new_session_builder, unique_keyspace_name}, + Session, +}; +use std::collections::BTreeSet; + +#[tokio::test] +async fn test_quietly_prepare_batch() { + let session = create_new_session_builder().build().await.unwrap(); + + let ks = unique_keyspace_name(); + session.query(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap(); + session.use_keyspace(ks.clone(), false).await.unwrap(); + + session + .query( + "CREATE TABLE 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 unprepared_insert_1_2: &str = "insert into test_batch_table (a, b) values (1, 2)"; + let unprepared_insert_2_3: &str = "insert into test_batch_table (a, b) values (2, 3)"; + let unprepared_insert_3_4: &str = "insert into test_batch_table (a, b) values (3, 4)"; + let unprepared_insert_4_5: &str = "insert into test_batch_table (a, b) values (4, 5)"; + let prepared_insert_a_b: PreparedStatement = + session.prepare(unprepared_insert_a_b).await.unwrap(); + let prepared_insert_a_7: PreparedStatement = + session.prepare(unprepared_insert_a_7).await.unwrap(); + let prepared_insert_8_b: PreparedStatement = + session.prepare(unprepared_insert_8_b).await.unwrap(); + + { + 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.clone()); + 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 mut unprepared_batch1: Batch = Default::default(); + unprepared_batch1.append_statement(unprepared_insert_1_2); + unprepared_batch1.append_statement(unprepared_insert_2_3); + unprepared_batch1.append_statement(unprepared_insert_3_4); + + session + .batch(&unprepared_batch1, ((), (), ())) + .await + .unwrap(); + assert_test_batch_table_rows_contain(&session, &[(1, 2), (2, 3), (3, 4)]).await; + } + + { + let mut unprepared_batch2: Batch = Default::default(); + unprepared_batch2.append_statement(unprepared_insert_a_b); + unprepared_batch2.append_statement(unprepared_insert_a_7); + unprepared_batch2.append_statement(unprepared_insert_8_b); + + session + .batch(&unprepared_batch2, ((12, 22), (12,), (22,))) + .await + .unwrap(); + assert_test_batch_table_rows_contain(&session, &[(12, 22), (12, 7), (8, 22)]).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); + partially_prepared_batch.append_statement(unprepared_insert_4_5); + + session + .batch(&partially_prepared_batch, ((33, 43), (33,), ())) + .await + .unwrap(); + assert_test_batch_table_rows_contain(&session, &[(33, 43), (33, 7), (4, 5)]).await; + } +} + +async fn assert_test_batch_table_rows_contain(sess: &Session, expected_rows: &[(i32, i32)]) { + let selected_rows: BTreeSet<(i32, i32)> = sess + .query("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 + ); + } + } +}