From fc35fa9ac9f231ae50ce4b53805b6184481e6907 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?David=20L=C3=B6nnhager?= <david.l@mullvad.net>
Date: Tue, 9 Apr 2024 15:24:29 +0200
Subject: [PATCH] test: use repeated probes

---
 test/test-manager/src/tests/helpers.rs  | 29 ++++++++++++++++++++++++-
 test/test-manager/src/tests/settings.rs | 18 +++++----------
 2 files changed, 33 insertions(+), 14 deletions(-)

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index bd24770fe7e6..10ea343f4f09 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -55,7 +55,7 @@ pub async fn reboot(rpc: &mut ServiceClient) -> Result<(), Error> {
     Ok(())
 }
 
-#[derive(Debug, Default)]
+#[derive(Debug, Default, PartialEq)]
 pub struct ProbeResult {
     tcp: usize,
     udp: usize,
@@ -95,6 +95,33 @@ pub async fn get_tunnel_interface(client: &mut MullvadProxyClient) -> Option<Str
     }
 }
 
+/// Same as send_guest_probes, but repeated until a consistent result is observed (but giving up
+/// after a number of attempts). This is useful where we're not necessarily in the desired state
+/// yet.
+pub async fn send_repeated_guest_probes(
+    rpc: ServiceClient,
+    interface: String,
+    destination: SocketAddr,
+) -> ProbeResult {
+    const MAX_ATTEMPTS: usize = 5;
+    const RETRY_INTERVAL: Duration = Duration::from_secs(3);
+
+    let previous_result = None;
+
+    for _ in 0..MAX_ATTEMPTS {
+        let new_result = send_guest_probes(rpc.clone(), interface.clone(), destination).await;
+
+        if previous_result == Some(new_result) {
+            return new_result;
+        }
+
+        tokio::time::sleep(RETRY_INTERVAL).await;
+    }
+
+    // cannot panic as we've received at least one result
+    previous_result.unwrap()
+}
+
 /// Sends a number of probes and returns the number of observed packets (UDP, TCP, or ICMP)
 pub async fn send_guest_probes(
     rpc: ServiceClient,
diff --git a/test/test-manager/src/tests/settings.rs b/test/test-manager/src/tests/settings.rs
index 01325ef32bb1..bb8ba6ea5d24 100644
--- a/test/test-manager/src/tests/settings.rs
+++ b/test/test-manager/src/tests/settings.rs
@@ -1,6 +1,6 @@
 use super::{
     helpers,
-    helpers::{connect_and_wait, send_guest_probes},
+    helpers::{connect_and_wait, send_guest_probes, send_repeated_guest_probes},
     Error, TestContext,
 };
 use crate::vm::network::DUMMY_LAN_INTERFACE_IP;
@@ -61,15 +61,13 @@ pub async fn test_lan(
         .await
         .expect("failed to enable LAN sharing");
 
-    // Firewall updates are a bit racey, so wait for a while
-    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
-
     // Ensure LAN is reachable
     //
 
     log::info!("Test whether outgoing LAN traffic is blocked");
 
-    let detected_probes = send_guest_probes(rpc.clone(), default_interface, lan_destination).await;
+    let detected_probes =
+        send_repeated_guest_probes(rpc.clone(), default_interface, lan_destination).await;
     assert!(
         detected_probes.all(),
         "did not observe all outgoing LAN packets: {detected_probes:?}"
@@ -117,16 +115,13 @@ pub async fn test_lockdown(
         .await
         .expect("failed to enable lockdown mode");
 
-    // Firewall updates are a bit racey, so wait for a while
-    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
-
     // Ensure all destinations are unreachable
     //
 
     let default_interface = rpc.get_default_interface().await?;
 
     let detected_probes =
-        send_guest_probes(rpc.clone(), default_interface.clone(), lan_destination).await;
+        send_repeated_guest_probes(rpc.clone(), default_interface.clone(), lan_destination).await;
     assert!(
         detected_probes.none(),
         "observed outgoing packets to LAN: {detected_probes:?}"
@@ -149,14 +144,11 @@ pub async fn test_lockdown(
         .await
         .expect("failed to enable LAN sharing");
 
-    // Firewall updates are a bit racey, so wait for a while
-    tokio::time::sleep(std::time::Duration::from_secs(3)).await;
-
     // Ensure private IPs are reachable, but not others
     //
 
     let detected_probes =
-        send_guest_probes(rpc.clone(), default_interface.clone(), lan_destination).await;
+        send_repeated_guest_probes(rpc.clone(), default_interface.clone(), lan_destination).await;
     assert!(
         detected_probes.all(),
         "did not observe some outgoing packets: {detected_probes:?}"