Fix wireguard rotation test

The test was flaky because if a race condition which made the key
rotation missable.

2 files changed, 41 insertions, 22 deletions

diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index 847150c477..fd6c56e19f 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -16,8 +16,7 @@ use mullvad_types::{
     version::AppVersionInfo,
     wireguard::{PublicKey, QuantumResistantState, RotationInterval},
 };
-use std::path::Path;
-use std::str::FromStr;
+use std::{path::Path, str::FromStr};
 #[cfg(target_os = "windows")]
 use talpid_types::split_tunnel::ExcludedProcess;
 use tonic::{Code, Status};
@@ -29,6 +28,7 @@ pub type Result<T> = std::result::Result<T, super::Error>;
 #[derive(Debug, Clone)]
 pub struct MullvadProxyClient(crate::ManagementServiceClient);
 
+#[derive(Debug)]
 pub enum DaemonEvent {
     TunnelState(TunnelState),
     Settings(Settings),
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(())
 }