Fix flaky kv test

Signed-off-by: Tomasz Pietrek <[email protected]>
nats-io · Mar 4, 2024 · 22a9718 · 22a9718
1 parent db70181
commit 22a9718
Showing 1 changed file with 16 additions and 11 deletions.
diff --git a/async-nats/tests/kv_tests.rs b/async-nats/tests/kv_tests.rs
@@ -834,17 +834,22 @@ mod kv {
             .unwrap();
 
         kv.put("baz", "value".into()).await.unwrap();
-        tokio::time::sleep(Duration::from_millis(300)).await;
-        assert_eq!(
-            kv.keys()
-                .await
-                .unwrap()
-                .try_collect::<Vec<String>>()
-                .await
-                .unwrap()
-                .len(),
-            0
-        );
+        tryhard::retry_fn(|| async {
+            match kv.keys().await {
+                Ok(keys) => {
+                    let keys = keys.try_collect::<Vec<String>>().await.unwrap();
+                    if !keys.is_empty() {
+                        return Err("keys not empty".into());
+                    }
+                    Ok::<(), async_nats::Error>(())
+                }
+                Err(e) => Err(e.into()),
+            }
+        })
+        .retries(5)
+        .exponential_backoff(Duration::from_millis(100))
+        .await
+        .unwrap();
     }
 
     #[tokio::test]