review feedback: add request timeout to execution profile to enable t…

…hat writes complete during cassandra tests

scylla/src/transport/large_batch_statements_test.rs

@@ -3,27 +3,34 @@ use crate::query::Query;
 use crate::{
     batch::Batch,
     test_utils::{create_new_session_builder, unique_keyspace_name},
-    QueryResult, Session,
+    ExecutionProfile, QueryResult, Session,
 };
 use assert_matches::assert_matches;
-use scylla_cql::errors::{BadQuery, DbError, QueryError};
+use scylla_cql::errors::{BadQuery, QueryError};
+use std::time::Duration;
 
 #[tokio::test]
 async fn test_large_batch_statements() {
-    let mut session = create_new_session_builder().build().await.unwrap();
+    let profile = ExecutionProfile::builder()
+        // set a 1 minute request timeout which should be enough to write this batch statement
+        .request_timeout(Some(Duration::from_secs(60)))
+        .build();
+
+    let handle = profile.into_handle();
+    let mut session = create_new_session_builder()
+        .default_execution_profile_handle(handle)
+        .build()
+        .await
+        .unwrap();
+
     let ks = unique_keyspace_name();
     session = create_test_session(session, &ks).await;
 
     // Add batch
     let max_number_of_queries = u16::MAX as usize;
     let batch_result = write_batch(&session, max_number_of_queries, &ks).await;
 
-    if batch_result.is_err() {
-        assert_matches!(
-            batch_result.unwrap_err(),
-            QueryError::DbError(DbError::WriteTimeout { .. }, _)
-        )
-    }
+    assert!(batch_result.is_ok());
 
     // Now try with too many queries
     let too_many_queries = u16::MAX as usize + 1;