Reduce wait time in integration test.

Poll for server and connect when ready.

BUG: 358538612
Change-Id: I5543aa4cce36c4f30c146de757747764122ef9db

oak_functions/examples/key_value_lookup/module/src/tests.rs

@@ -19,7 +19,7 @@
 
 extern crate test;
 
-use std::time::Duration;
+use std::time::{Duration, Instant};
 
 use maplit::hashmap;
 use oak_client::verifier::InsecureAttestationVerifier;
@@ -100,13 +100,26 @@ fn bench_wasm_handler(bencher: &mut Bencher) {
         lookup_data_file.path().to_str().unwrap(),
     );
 
-    // Wait for the server to start up.
-    std::thread::sleep(Duration::from_secs(180));
-
     let uri = format!("http://localhost:{server_port}/");
-    let mut client = runtime
-        .block_on(OakFunctionsClient::new(&uri, &InsecureAttestationVerifier))
-        .expect("couldn't create client");
+
+    // Try to connect while waiting for the server to start up.
+    let timeout = Duration::from_secs(180);
+    let start = Instant::now();
+
+    let mut client = loop {
+        log::debug!("Waiting for server at {}...", uri);
+        let maybe_client =
+            runtime.block_on(OakFunctionsClient::new(&uri, &InsecureAttestationVerifier));
+        if maybe_client.is_ok() {
+            log::debug!("Connected!");
+            break maybe_client.unwrap();
+        } else {
+            if (Instant::now() - start) > timeout {
+                maybe_client.expect("couldn't create client");
+            }
+            std::thread::sleep(Duration::from_secs(5));
+        }
+    };
 
     let summary = bencher.bench(|bencher| {
         bencher.iter(|| {