From abfb81093d744a0185cf39054c99921caab07992 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Karol=20Bary=C5=82a?= <karol.baryla@scylladb.com>
Date: Mon, 14 Oct 2024 12:03:44 +0200
Subject: [PATCH] default policy: Move some tests to integration

---
 .../src/transport/load_balancing/default.rs   | 25 +-----------------
 scylla/tests/integration/default_policy.rs    | 26 +++++++++++++++++++
 scylla/tests/integration/main.rs              |  1 +
 3 files changed, 28 insertions(+), 24 deletions(-)
 create mode 100644 scylla/tests/integration/default_policy.rs

diff --git a/scylla/src/transport/load_balancing/default.rs b/scylla/src/transport/load_balancing/default.rs
index 51db7f97f..8e35cd133 100644
--- a/scylla/src/transport/load_balancing/default.rs
+++ b/scylla/src/transport/load_balancing/default.rs
@@ -3126,7 +3126,7 @@ mod latency_awareness {
         use crate::{
             load_balancing::default::NodeLocationPreference,
             routing::Shard,
-            test_utils::{create_new_session_builder, setup_tracing},
+            test_utils::setup_tracing,
             transport::locator::test::{TABLE_INVALID, TABLE_NTS_RF_2, TABLE_NTS_RF_3},
         };
         use crate::{
@@ -3139,7 +3139,6 @@ mod latency_awareness {
                 locator::test::{id_to_invalid_addr, A, B, C, D, E, F, G},
                 ClusterData, NodeAddr,
             },
-            ExecutionProfile,
         };
         use tokio::time::Instant;
 
@@ -3845,28 +3844,6 @@ mod latency_awareness {
             }
         }
 
-        // This is a regression test for #696.
-        #[tokio::test]
-        #[ntest::timeout(1000)]
-        async fn latency_aware_query_completes() {
-            setup_tracing();
-            let policy = DefaultPolicy::builder()
-                .latency_awareness(LatencyAwarenessBuilder::default())
-                .build();
-            let handle = ExecutionProfile::builder()
-                .load_balancing_policy(policy)
-                .build()
-                .into_handle();
-
-            let session = create_new_session_builder()
-                .default_execution_profile_handle(handle)
-                .build()
-                .await
-                .unwrap();
-
-            session.query_unpaged("whatever", ()).await.unwrap_err();
-        }
-
         #[tokio::test(start_paused = true)]
         async fn timestamped_average_works_when_clock_stops() {
             setup_tracing();
diff --git a/scylla/tests/integration/default_policy.rs b/scylla/tests/integration/default_policy.rs
new file mode 100644
index 000000000..e097cf418
--- /dev/null
+++ b/scylla/tests/integration/default_policy.rs
@@ -0,0 +1,26 @@
+use scylla::load_balancing::{DefaultPolicy, LatencyAwarenessBuilder};
+use scylla::ExecutionProfile;
+
+use crate::utils::{create_new_session_builder, setup_tracing};
+
+// This is a regression test for #696.
+#[tokio::test]
+#[ntest::timeout(1000)]
+async fn latency_aware_query_completes() {
+    setup_tracing();
+    let policy = DefaultPolicy::builder()
+        .latency_awareness(LatencyAwarenessBuilder::default())
+        .build();
+    let handle = ExecutionProfile::builder()
+        .load_balancing_policy(policy)
+        .build()
+        .into_handle();
+
+    let session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
+    session.query_unpaged("whatever", ()).await.unwrap_err();
+}
diff --git a/scylla/tests/integration/main.rs b/scylla/tests/integration/main.rs
index 24dfe731a..d510dc6fd 100644
--- a/scylla/tests/integration/main.rs
+++ b/scylla/tests/integration/main.rs
@@ -3,6 +3,7 @@ mod consistency;
 mod cql_collections;
 mod cql_types;
 mod cql_value;
+mod default_policy;
 mod execution_profiles;
 mod history;
 mod hygiene;