From f87e1338de214b4f90ff388e2da0ec804e32ad53 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Karol=20Bary=C5=82a?= <karol.baryla@scylladb.com>
Date: Wed, 20 Nov 2024 15:55:41 +0100
Subject: [PATCH] Tests: Change the way tablet support is checked

Fixes: https://github.com/scylladb/scylla-rust-driver/issues/1048

See the issue description for more information.
---
 scylla/src/utils/test_utils.rs    | 15 +--------------
 scylla/tests/integration/utils.rs | 15 +--------------
 2 files changed, 2 insertions(+), 28 deletions(-)

diff --git a/scylla/src/utils/test_utils.rs b/scylla/src/utils/test_utils.rs
index 68035fcee..d745a3135 100644
--- a/scylla/src/utils/test_utils.rs
+++ b/scylla/src/utils/test_utils.rs
@@ -1,6 +1,5 @@
 use crate::transport::session_builder::{GenericSessionBuilder, SessionBuilderKind};
 use crate::Session;
-use scylla_cql::frame::response::result::Row;
 use std::{num::NonZeroU32, time::Duration};
 use std::{
     sync::atomic::{AtomicUsize, Ordering},
@@ -92,19 +91,7 @@ pub(crate) fn create_new_session_builder() -> GenericSessionBuilder<impl Session
 }
 
 pub(crate) async fn scylla_supports_tablets(session: &Session) -> bool {
-    let result = session
-        .query_unpaged(
-            "select column_name from system_schema.columns where
-                keyspace_name = 'system_schema'
-                and table_name = 'scylla_keyspaces'
-                and column_name = 'initial_tablets'",
-            &[],
-        )
-        .await
-        .unwrap()
-        .into_rows_result();
-
-    result.is_ok_and(|rows_result| rows_result.single_row::<Row>().is_ok())
+    supports_feature(session, "TABLETS").await
 }
 
 pub(crate) fn setup_tracing() {
diff --git a/scylla/tests/integration/utils.rs b/scylla/tests/integration/utils.rs
index 4ec8539c5..ecbffcb1d 100644
--- a/scylla/tests/integration/utils.rs
+++ b/scylla/tests/integration/utils.rs
@@ -1,6 +1,5 @@
 use futures::Future;
 use scylla::deserialize::DeserializeValue;
-use scylla::frame::response::result::Row;
 use scylla::transport::session_builder::{GenericSessionBuilder, SessionBuilderKind};
 use scylla::Session;
 use std::collections::HashMap;
@@ -123,19 +122,7 @@ pub(crate) async fn supports_feature(session: &Session, feature: &str) -> bool {
 
 #[allow(unused)]
 pub(crate) async fn scylla_supports_tablets(session: &Session) -> bool {
-    let result = session
-        .query_unpaged(
-            "select column_name from system_schema.columns where
-                keyspace_name = 'system_schema'
-                and table_name = 'scylla_keyspaces'
-                and column_name = 'initial_tablets'",
-            &[],
-        )
-        .await
-        .unwrap()
-        .into_rows_result();
-
-    result.is_ok_and(|rows_result| rows_result.single_row::<Row>().is_ok())
+    supports_feature(session, "TABLETS").await
 }
 
 // Creates a generic session builder based on conditional compilation configuration