Move timestamp test to statement.rs and add missing cases

Author: sylwiaszunejko

scylla/tests/integration/main.rs

@@ -19,6 +19,7 @@ mod shards;
 mod silent_prepare_batch;
 mod silent_prepare_query;
 mod skip_metadata_optimization;
+mod statement;
 mod tablets;
 #[path = "../common/utils.rs"]
 mod utils;

scylla/tests/integration/session.rs

@@ -5,7 +5,7 @@ use crate::utils::{
 use assert_matches::assert_matches;
 use futures::{FutureExt, StreamExt as _, TryStreamExt};
 use itertools::Itertools;
-use scylla::batch::{Batch, BatchStatement, BatchType};
+use scylla::batch::{Batch, BatchStatement};
 use scylla::client::caching_session::CachingSession;
 use scylla::client::execution_profile::ExecutionProfile;
 use scylla::client::session::Session;
@@ -31,7 +31,7 @@ use std::collections::BTreeSet;
 use std::collections::{BTreeMap, HashMap};
 use std::collections::HashSet;
 use std::sync::atomic::{AtomicBool, Ordering};
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
 use tokio::net::TcpListener;
 use uuid::Uuid;
 
@@ -945,198 +945,6 @@ async fn test_await_schema_agreement() {
     let _schema_version = session.await_schema_agreement().await.unwrap();
 }
 
-#[tokio::test]
-async fn test_timestamp() {
-    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
-        .ddl(format!(
-            "CREATE TABLE IF NOT EXISTS {}.t_timestamp (a text, b text, primary key (a))",
-            ks
-        ))
-        .await
-        .unwrap();
-
-    session.await_schema_agreement().await.unwrap();
-
-    let query_str = format!("INSERT INTO {}.t_timestamp (a, b) VALUES (?, ?)", ks);
-
-    // test regular query timestamps
-
-    let mut regular_query = Query::new(query_str.to_string());
-
-    regular_query.set_timestamp(Some(420));
-    session
-        .query_unpaged(regular_query.clone(), ("regular query", "higher timestamp"))
-        .await
-        .unwrap();
-
-    regular_query.set_timestamp(Some(42));
-    session
-        .query_unpaged(regular_query.clone(), ("regular query", "lower timestamp"))
-        .await
-        .unwrap();
-
-    // test prepared statement timestamps
-
-    let mut prepared_statement = session.prepare(query_str).await.unwrap();
-
-    prepared_statement.set_timestamp(Some(420));
-    session
-        .execute_unpaged(&prepared_statement, ("prepared query", "higher timestamp"))
-        .await
-        .unwrap();
-
-    prepared_statement.set_timestamp(Some(42));
-    session
-        .execute_unpaged(&prepared_statement, ("prepared query", "lower timestamp"))
-        .await
-        .unwrap();
-
-    // test batch statement timestamps
-
-    let mut batch: Batch = Default::default();
-    batch.append_statement(regular_query);
-    batch.append_statement(prepared_statement);
-
-    batch.set_timestamp(Some(420));
-    session
-        .batch(
-            &batch,
-            (
-                ("first query in batch", "higher timestamp"),
-                ("second query in batch", "higher timestamp"),
-            ),
-        )
-        .await
-        .unwrap();
-
-    batch.set_timestamp(Some(42));
-    session
-        .batch(
-            &batch,
-            (
-                ("first query in batch", "lower timestamp"),
-                ("second query in batch", "lower timestamp"),
-            ),
-        )
-        .await
-        .unwrap();
-
-    let query_rows_result = session
-        .query_unpaged(
-            format!("SELECT a, b, WRITETIME(b) FROM {}.t_timestamp", ks),
-            &[],
-        )
-        .await
-        .unwrap()
-        .into_rows_result()
-        .unwrap();
-
-    let mut results = query_rows_result
-        .rows::<(&str, &str, i64)>()
-        .unwrap()
-        .map(Result::unwrap)
-        .collect::<Vec<_>>();
-    results.sort();
-
-    let expected_results = [
-        ("first query in batch", "higher timestamp", 420),
-        ("prepared query", "higher timestamp", 420),
-        ("regular query", "higher timestamp", 420),
-        ("second query in batch", "higher timestamp", 420),
-    ]
-    .into_iter()
-    .collect::<Vec<_>>();
-
-    assert_eq!(results, expected_results);
-}
-
-#[tokio::test]
-async fn test_timestamp_generator() {
-    use rand::random;
-    use scylla::policies::timestamp_generator::TimestampGenerator;
-
-    setup_tracing();
-    struct LocalTimestampGenerator {
-        generated_timestamps: Arc<Mutex<HashSet<i64>>>,
-    }
-
-    impl TimestampGenerator for LocalTimestampGenerator {
-        fn next_timestamp(&self) -> i64 {
-            let timestamp = random::<i64>().abs();
-            self.generated_timestamps.lock().unwrap().insert(timestamp);
-            timestamp
-        }
-    }
-
-    let timestamps = Arc::new(Mutex::new(HashSet::new()));
-    let generator = LocalTimestampGenerator {
-        generated_timestamps: timestamps.clone(),
-    };
-
-    let session = create_new_session_builder()
-        .timestamp_generator(Arc::new(generator))
-        .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
-        .ddl(format!(
-            "CREATE TABLE IF NOT EXISTS {}.t_generator (a int primary key, b int)",
-            ks
-        ))
-        .await
-        .unwrap();
-
-    let prepared = session
-        .prepare(format!(
-            "INSERT INTO {}.t_generator (a, b) VALUES (1, 1)",
-            ks
-        ))
-        .await
-        .unwrap();
-    session.execute_unpaged(&prepared, []).await.unwrap();
-
-    let unprepared = Query::new(format!(
-        "INSERT INTO {}.t_generator (a, b) VALUES (2, 2)",
-        ks
-    ));
-    session.query_unpaged(unprepared, []).await.unwrap();
-
-    let mut batch = Batch::new(BatchType::Unlogged);
-    let stmt = session
-        .prepare(format!(
-            "INSERT INTO {}.t_generator (a, b) VALUES (3, 3)",
-            ks
-        ))
-        .await
-        .unwrap();
-    batch.append_statement(stmt);
-    session.batch(&batch, &((),)).await.unwrap();
-
-    let query_rows_result = session
-        .query_unpaged(
-            format!("SELECT a, b, WRITETIME(b) FROM {}.t_generator", ks),
-            &[],
-        )
-        .await
-        .unwrap()
-        .into_rows_result()
-        .unwrap();
-
-    let timestamps_locked = timestamps.lock().unwrap();
-    assert!(query_rows_result
-        .rows::<(i32, i32, i64)>()
-        .unwrap()
-        .map(|row_result| row_result.unwrap())
-        .all(|(_a, _b, writetime)| timestamps_locked.contains(&writetime)));
-}
-
 #[ignore = "works on remote Scylla instances only (local ones are too fast)"]
 #[tokio::test]
 async fn test_request_timeout() {