From 6d1943a78c72e9890376473d19ee0be8fdf94b55 Mon Sep 17 00:00:00 2001
From: smoczy123 <michal.hudobski@poczta.onet.pl>
Date: Thu, 5 Dec 2024 01:01:08 +0100
Subject: [PATCH] transport/session_test: Added unit and integration test for
 MonotonicTimestampGenerator

---
 scylla/src/transport/session_test.rs        | 59 +++++++++++++++++++++
 scylla/src/transport/timestamp_generator.rs | 16 ++++++
 2 files changed, 75 insertions(+)

diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index f3bb159bc..06afdc2ad 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -1307,6 +1307,65 @@ async fn test_timestamp() {
     assert_eq!(results, expected_results);
 }
 
+#[tokio::test]
+async fn test_timestamp_generator() {
+    setup_tracing();
+    use crate::transport::timestamp_generator::MonotonicTimestampGenerator;
+
+    let session = create_new_session_builder()
+        .timestamp_generator(Arc::new(MonotonicTimestampGenerator::new()))
+        .build()
+        .await
+        .unwrap();
+    let ks = unique_keyspace_name();
+
+    session.query_unpaged(format!("CREATE KEYSPACE IF NOT EXISTS {} WITH REPLICATION = {{'class' : 'NetworkTopologyStrategy', 'replication_factor' : 1}}", ks), &[]).await.unwrap();
+    session
+        .query_unpaged(
+            format!(
+                "CREATE TABLE IF NOT EXISTS {}.t_generator (a int, b int, primary key (a))",
+                ks
+            ),
+            &[],
+        )
+        .await
+        .unwrap();
+
+    session.await_schema_agreement().await.unwrap();
+
+    for n in 1..100 {
+        let prepared = session
+            .prepare(format!(
+                "INSERT INTO {}.t_generator (a, b) VALUES (?, ?)",
+                ks
+            ))
+            .await
+            .unwrap();
+        session.execute_unpaged(&prepared, (n, n)).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 mut results = query_rows_result
+        .rows::<(i32, i32, i64)>()
+        .unwrap()
+        .map(Result::unwrap)
+        .collect::<Vec<_>>();
+    results.sort();
+    let timestamps = results.into_iter().map(|x| x.2).collect::<Vec<_>>();
+    let mut sorted_timestamps = timestamps.clone();
+    sorted_timestamps.sort();
+    assert_eq!(timestamps, sorted_timestamps);
+}
+
 #[ignore = "works on remote Scylla instances only (local ones are too fast)"]
 #[tokio::test]
 async fn test_request_timeout() {
diff --git a/scylla/src/transport/timestamp_generator.rs b/scylla/src/transport/timestamp_generator.rs
index 99eb7372f..bd736cf3e 100644
--- a/scylla/src/transport/timestamp_generator.rs
+++ b/scylla/src/transport/timestamp_generator.rs
@@ -93,3 +93,19 @@ impl TimestampGenerator for MonotonicTimestampGenerator {
         }
     }
 }
+
+#[tokio::test]
+async fn monotonic_timestamp_generator_is_monotonic() {
+    const NUMBER_OF_ITERATIONS: u32 = 1000;
+
+    let mut prev = None;
+    let mut cur;
+    let generator = MonotonicTimestampGenerator::new();
+    for _ in 0..NUMBER_OF_ITERATIONS {
+        cur = generator.next_timestamp().await;
+        if let Some(prev_val) = prev {
+            assert!(cur > prev_val);
+        }
+        prev = Some(cur);
+    }
+}