diff --git a/src/auth/src/user_provider/watch_file_user_provider.rs b/src/auth/src/user_provider/watch_file_user_provider.rs index 19b6530f21a3..0ca6ae09386b 100644 --- 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, ) { - 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,10 +174,11 @@ 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(); @@ -169,19 +186,17 @@ pub mod test { 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; } }