From 01fe25fe9c087e555f2716b1a9157c22868094aa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Wojciech=20Przytu=C5=82a?= <wojciech.przytula@scylladb.com>
Date: Mon, 22 Aug 2022 10:04:28 +0200
Subject: [PATCH] session_test: added test for request timeouts

As tokio timeout supports minimal granularity of 1 ms, the test can
be only reliable on remote Scyllas, as local ones respond faster.
Therefore, the rest is marked as ignored.
As all ignored tests were run in Authenticate CI workflow, it was made
specific to avoid running the new request timeouts test.
---
 .github/workflows/authenticate_test.yml |  2 +-
 scylla/Cargo.toml                       |  1 +
 scylla/src/transport/session_test.rs    | 69 +++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/authenticate_test.yml b/.github/workflows/authenticate_test.yml
index 620f7e9e3d..5f2b271d67 100644
--- a/.github/workflows/authenticate_test.yml
+++ b/.github/workflows/authenticate_test.yml
@@ -24,4 +24,4 @@ jobs:
     steps:
     - uses: actions/checkout@v2
     - name: Run tests
-      run: cargo test --verbose --  --ignored
+      run: cargo test --verbose authenticate_superuser --  --ignored
diff --git a/scylla/Cargo.toml b/scylla/Cargo.toml
index 52fcf7b86a..c98b9dccc0 100644
--- a/scylla/Cargo.toml
+++ b/scylla/Cargo.toml
@@ -47,6 +47,7 @@ smallvec = "1.8.0"
 [dev-dependencies]
 criterion = "0.3"
 tracing-subscriber = "0.3.14"
+assert_matches = "1.5.0"
 
 [[bench]]
 name = "benchmark"
diff --git a/scylla/src/transport/session_test.rs b/scylla/src/transport/session_test.rs
index 5c1248ee95..f95d8401f2 100644
--- a/scylla/src/transport/session_test.rs
+++ b/scylla/src/transport/session_test.rs
@@ -14,6 +14,7 @@ use crate::transport::topology::{CollectionType, ColumnKind, CqlType, NativeType
 use crate::CachingSession;
 use crate::QueryResult;
 use crate::{IntoTypedRows, Session, SessionBuilder};
+use assert_matches::assert_matches;
 use bytes::Bytes;
 use futures::{FutureExt, StreamExt};
 use itertools::Itertools;
@@ -1181,6 +1182,74 @@ async fn test_timestamp() {
     assert_eq!(results, expected_results);
 }
 
+#[ignore = "works on remote Scylla instances only (local ones are too fast)"]
+#[tokio::test]
+async fn test_request_timeout() {
+    use std::time::Duration;
+    let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());
+
+    {
+        let session = SessionBuilder::new()
+            .known_node(uri.as_str())
+            .build()
+            .await
+            .unwrap();
+
+        let mut query: Query = Query::new("SELECT * FROM system_schema.tables");
+        query.set_request_timeout(Some(Duration::from_millis(1)));
+        match session.query(query, &[]).await {
+            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        }
+
+        let mut prepared = session
+            .prepare("SELECT * FROM system_schema.tables")
+            .await
+            .unwrap();
+
+        prepared.set_request_timeout(Some(Duration::from_millis(1)));
+        match session.execute(&prepared, &[]).await {
+            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+    }
+    {
+        let timeouting_session = SessionBuilder::new()
+            .known_node(uri)
+            .request_timeout(Some(Duration::from_millis(1)))
+            .build()
+            .await
+            .unwrap();
+
+        let mut query = Query::new("SELECT * FROM system_schema.tables");
+
+        match timeouting_session.query(query.clone(), &[]).await {
+            Ok(_) => panic!("the query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+
+        query.set_request_timeout(Some(Duration::from_secs(10000)));
+
+        timeouting_session.query(query, &[]).await.expect(
+            "the query should have not failed, because no client-side timeout was specified",
+        );
+
+        let mut prepared = timeouting_session
+            .prepare("SELECT * FROM system_schema.tables")
+            .await
+            .unwrap();
+
+        match timeouting_session.execute(&prepared, &[]).await {
+            Ok(_) => panic!("the prepared query should have failed due to a client-side timeout"),
+            Err(e) => assert_matches!(e, QueryError::RequestTimeout(_)),
+        };
+
+        prepared.set_request_timeout(Some(Duration::from_secs(10000)));
+
+        timeouting_session.execute(&prepared, &[]).await.expect("the prepared query should have not failed, because no client-side timeout was specified");
+    }
+}
+
 #[tokio::test]
 async fn test_prepared_config() {
     let uri = std::env::var("SCYLLA_URI").unwrap_or_else(|_| "127.0.0.1:9042".to_string());