1 files changed, 39 insertions, 20 deletions

diff --git a/test/test-manager/src/tests/account.rs b/test/test-manager/src/tests/account.rs
index 94374dba20..69f605a290 100644
--- a/test/test-manager/src/tests/account.rs
+++ b/test/test-manager/src/tests/account.rs
@@ -288,50 +288,69 @@ pub async fn test_automatic_wireguard_rotation(
         .device
         .pubkey;
 
-    // Stop daemon
+    log::info!("Old wireguard key: {old_key}");
+
+    log::info!("Stopping daemon");
     rpc.stop_mullvad_daemon()
         .await
         .expect("Could not stop system service");
 
-    // Open device.json and change created field to more than 7 days ago
+    log::info!("Changing created field of `device.json` to more than 7 days ago");
     rpc.make_device_json_old()
         .await
         .expect("Could not change device.json to have an old created timestamp");
 
-    // Start daemon
+    log::info!("Starting daemon");
     rpc.start_mullvad_daemon()
         .await
         .expect("Could not start system service");
 
     // NOTE: Need to create a new `mullvad_client` here after the restart otherwise we can't
     // communicate with the daemon
+    log::info!("Reconnecting to daemon");
     drop(mullvad_client);
     let mut mullvad_client = ctx.rpc_provider.new_client().await;
 
-    // Verify rotation has happened after a minute
-    const KEY_ROTATION_TIMEOUT: Duration = Duration::from_secs(100);
+    log::info!("Verifying that wireguard key has change");
+
+    // Check if the key rotation has already occurred when connected to the daemon, otherwise
+    // listen for device daemon events until we observe the change. We have to register the event
+    // listener before polling the current key to be sure we don't miss the change.
+    let event_listener = mullvad_client.events_listen().await.unwrap();
+    let new_key = mullvad_client
+        .get_device()
+        .await
+        .unwrap()
+        .into_device()
+        .expect("Could not get device")
+        .device
+        .pubkey;
 
-    let new_key = tokio::time::timeout(
-        KEY_ROTATION_TIMEOUT,
-        helpers::find_daemon_event(
-            mullvad_client.events_listen().await.unwrap(),
-            |daemon_event| match daemon_event {
+    // If key has not yet been updated, listen for changes to it
+    if new_key == old_key {
+        log::info!("Listening for device daemon event");
+        // Verify rotation has happened within 100 seconds - if the key hasn't been rotated after
+        // that, the rotation probably won't happen anytime soon.
+        let device_event = tokio::task::spawn(tokio::time::timeout(
+            Duration::from_secs(100),
+            helpers::find_daemon_event(event_listener, |daemon_event| match daemon_event {
                 DaemonEvent::Device(device_event) => Some(device_event),
                 _ => None,
-            },
-        ),
-    )
-    .await
-    .map_err(|_error| Error::Daemon(String::from("Tunnel event listener timed out")))?
-    .map(|device_event| {
-        device_event
+            }),
+        ))
+        .await
+        .unwrap()
+        .map_err(|_error| Error::Daemon(String::from("Tunnel event listener timed out")))??;
+
+        let new_key = device_event
             .new_state
             .into_device()
             .expect("Could not get device")
             .device
-            .pubkey
-    })?;
+            .pubkey;
+
+        assert_ne!(old_key, new_key);
+    }
 
-    assert_ne!(old_key, new_key);
     Ok(())
 }