avoid sleep

Signed-off-by: tison <[email protected]>
GreptimeTeam · Mar 26, 2024 · 872b79d · 872b79d
1 parent b7f5336
commit 872b79d
Showing 1 changed file with 43 additions and 29 deletions.
diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs
@@ -114,7 +114,7 @@ impl UserProvider for WatchFileUserProvider {
 pub mod test {
     use std::fs::File;
     use std::io::{LineWriter, Write};
-    use std::time::Duration;
+    use std::time::{Duration, Instant};
 
     use common_test_util::temp_dir::create_temp_dir;
     use tokio::time::sleep;
@@ -128,20 +128,36 @@ pub mod test {
         username: &str,
         password: &str,
         ok: bool,
+        timeout: Option<Duration>,
     ) {
-        let re = provider
-            .authenticate(
-                Identity::UserId(username, None),
-                Password::PlainText(password.to_string().into()),
-            )
-            .await;
-        assert_eq!(
-            re.is_ok(),
-            ok,
-            "username: {}, password: {}",
-            username,
-            password
-        );
+        if let Some(timeout) = timeout {
+            let deadline = Instant::now().checked_add(timeout).unwrap();
+            loop {
+                let re = provider
+                    .authenticate(
+                        Identity::UserId(username, None),
+                        Password::PlainText(password.to_string().into()),
+                    )
+                    .await;
+                if re.is_err() && Instant::now() >= deadline {
+                    panic!("timeout (username: {username}, password: {password}, expected: {ok})");
+                }
+            }
+        } else {
+            let re = provider
+                .authenticate(
+                    Identity::UserId(username, None),
+                    Password::PlainText(password.to_string().into()),
+                )
+                .await;
+            assert_eq!(
+                re.is_ok(),
+                ok,
+                "username: {}, password: {}",
+                username,
+                password
+            );
+        }
     }
 
     #[tokio::test]
@@ -158,30 +174,29 @@ pub mod test {
         }
 
         let provider = WatchFileUserProvider::new(file_path.as_str()).unwrap();
-        test_authenticate(&provider, "root", "123456", true).await;
-        test_authenticate(&provider, "admin", "654321", true).await;
-        test_authenticate(&provider, "root", "654321", false).await;
+        test_authenticate(&provider, "root", "123456", true, None).await;
+        test_authenticate(&provider, "admin", "654321", true, None).await;
+        test_authenticate(&provider, "root", "654321", false, None).await;
 
+        let timeout = Duration::from_secs(2);
         {
             // update the tmp file
             let file = File::create(&file_path).unwrap();
             let mut lw = LineWriter::new(file);
             assert!(writeln!(lw, "root=654321").is_ok());
             lw.flush().unwrap();
         }
-        sleep(Duration::from_secs(2)).await; // wait the watcher to apply the change
-        test_authenticate(&provider, "root", "123456", false).await;
-        test_authenticate(&provider, "root", "654321", true).await;
-        test_authenticate(&provider, "admin", "654321", false).await;
+        test_authenticate(&provider, "root", "123456", false, Some(timeout)).await;
+        test_authenticate(&provider, "root", "654321", true, Some(timeout)).await;
+        test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await;
 
         {
             // remove the tmp file
             std::fs::remove_file(&file_path).unwrap();
         }
-        sleep(Duration::from_secs(2)).await; // wait the watcher to apply the change
-        test_authenticate(&provider, "root", "123456", true).await;
-        test_authenticate(&provider, "root", "654321", true).await;
-        test_authenticate(&provider, "admin", "654321", true).await;
+        test_authenticate(&provider, "root", "123456", true, Some(timeout)).await;
+        test_authenticate(&provider, "root", "654321", true, Some(timeout)).await;
+        test_authenticate(&provider, "admin", "654321", true, Some(timeout)).await;
 
         {
             // recreate the tmp file
@@ -190,9 +205,8 @@ pub mod test {
             assert!(writeln!(lw, "root=123456").is_ok());
             lw.flush().unwrap();
         }
-        sleep(Duration::from_secs(2)).await; // wait the watcher to apply the change
-        test_authenticate(&provider, "root", "123456", true).await;
-        test_authenticate(&provider, "root", "654321", false).await;
-        test_authenticate(&provider, "admin", "654321", false).await;
+        test_authenticate(&provider, "root", "123456", true, Some(timeout)).await;
+        test_authenticate(&provider, "root", "654321", false, Some(timeout)).await;
+        test_authenticate(&provider, "admin", "654321", false, Some(timeout)).await;
     }
 }