diff --git a/scylla/src/transport/large_batch_statements_test.rs b/scylla/src/transport/large_batch_statements_test.rs index b2a5e8eaf0..166ed0aaf0 100644 --- a/scylla/src/transport/large_batch_statements_test.rs +++ b/scylla/src/transport/large_batch_statements_test.rs @@ -3,14 +3,26 @@ 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 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; @@ -18,12 +30,7 @@ async fn test_large_batch_statements() { 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;