From 7e8c31c71597ef35eed0fae33773dacd10cdb205 Mon Sep 17 00:00:00 2001
From: barshaul <barshaul@amazon.com>
Date: Thu, 15 Aug 2024 07:51:11 +0000
Subject: [PATCH] Fixed tests

---
 redis/tests/test_cluster_async.rs |  6 +++---
 redis/tests/test_cluster_scan.rs  | 26 ++++++++++++++++++++++----
 2 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/redis/tests/test_cluster_async.rs b/redis/tests/test_cluster_async.rs
index 4dbc0653b..7d1249c3e 100644
--- a/redis/tests/test_cluster_async.rs
+++ b/redis/tests/test_cluster_async.rs
@@ -2400,7 +2400,7 @@ mod cluster_async {
             .block_on(cmd.query_async::<_, Vec<String>>(&mut connection))
             .unwrap_err();
         assert!(
-            matches!(result.kind(), ErrorKind::ClusterConnectionNotFound)
+            matches!(result.kind(), ErrorKind::ConnectionNotFoundForRoute)
                 || result.is_connection_dropped()
         );
     }
@@ -4031,7 +4031,7 @@ mod cluster_async {
             handler: _handler,
             ..
         } = MockEnv::with_client_builder(
-            ClusterClient::builder(vec![&*format!("redis://{name}")]),
+            ClusterClient::builder(vec![&*format!("redis://{name}")]).retries(1),
             name,
             move |received_cmd: &[u8], _| {
                 let slots_config_vec = vec![
@@ -4071,7 +4071,7 @@ mod cluster_async {
                 let res_err = res.unwrap_err();
                 assert_eq!(
                     res_err.kind(),
-                    ErrorKind::ClusterConnectionNotFound,
+                    ErrorKind::ConnectionNotFoundForRoute,
                     "{:?}",
                     res_err
                 );
diff --git a/redis/tests/test_cluster_scan.rs b/redis/tests/test_cluster_scan.rs
index a4bb85625..29a3c87b4 100644
--- a/redis/tests/test_cluster_scan.rs
+++ b/redis/tests/test_cluster_scan.rs
@@ -164,7 +164,12 @@ mod test_cluster_scan_async {
 
     #[tokio::test] // test cluster scan with node fail in the middle
     async fn test_async_cluster_scan_with_fail() {
-        let cluster = TestClusterContext::new(3, 0);
+        let cluster = TestClusterContext::new_with_cluster_client_builder(
+            3,
+            0,
+            |builder| builder.retries(1),
+            false,
+        );
         let mut connection = cluster.async_connection(None).await;
         // Set some keys
         for i in 0..1000 {
@@ -224,7 +229,11 @@ mod test_cluster_scan_async {
         let cluster = TestClusterContext::new_with_cluster_client_builder(
             6,
             1,
-            |builder| builder.slots_refresh_rate_limit(Duration::from_secs(0), 0),
+            |builder| {
+                builder
+                    .slots_refresh_rate_limit(Duration::from_secs(0), 0)
+                    .retries(1)
+            },
             false,
         );
 
@@ -374,7 +383,11 @@ mod test_cluster_scan_async {
         let cluster = TestClusterContext::new_with_cluster_client_builder(
             6,
             1,
-            |builder| builder.slots_refresh_rate_limit(Duration::from_secs(0), 0),
+            |builder| {
+                builder
+                    .slots_refresh_rate_limit(Duration::from_secs(0), 0)
+                    .retries(1)
+            },
             false,
         );
 
@@ -772,7 +785,12 @@ mod test_cluster_scan_async {
     // Testing cluster scan when connection fails in the middle and we get an error
     // then cluster up again and scanning can continue without any problem
     async fn test_async_cluster_scan_failover() {
-        let mut cluster = TestClusterContext::new(3, 0);
+        let mut cluster = TestClusterContext::new_with_cluster_client_builder(
+            3,
+            0,
+            |builder| builder.retries(1),
+            false,
+        );
         let mut connection = cluster.async_connection(None).await;
         let mut i = 0;
         loop {