From 9b06a90a2436f8777ddc570c8974bed6bf1fe6f5 Mon Sep 17 00:00:00 2001
From: Markus Pettersson <markus.pettersson@mullvad.net>
Date: Tue, 19 Mar 2024 10:41:35 +0100
Subject: [PATCH] Re-enable old test

---
 mullvad-daemon/src/device/mod.rs              | 36 ++++++++-----------
 .../src/relay_selector/mod.rs                 | 24 +++++++++++++
 2 files changed, 38 insertions(+), 22 deletions(-)

diff --git a/mullvad-daemon/src/device/mod.rs b/mullvad-daemon/src/device/mod.rs
index 67dc189bedcf..d8fcaa3903d6 100644
--- a/mullvad-daemon/src/device/mod.rs
+++ b/mullvad-daemon/src/device/mod.rs
@@ -1435,26 +1435,18 @@ mod test {
         );
     }
 
-    // TODO(markus): `preferred_tunnel_constraints` is slated for removal - consider writing a new test which
-    // does not depend on relay selector internals.
-    // /// Test whether the relay selector selects wireguard often enough, given no special
-    // /// constraints, to verify that the device is valid
-    // #[test]
-    // fn test_validates_by_default() {
-    //     for attempt in 0.. {
-    //         let should_validate =
-    //             TunnelStateChangeHandler::should_check_validity_on_attempt(attempt);
-    //         let (_, _, tunnel_type) =
-    //             RelaySelector::preferred_tunnel_constraints(attempt.try_into().unwrap());
-    //         assert_eq!(
-    //             tunnel_type,
-    //             TunnelType::Wireguard,
-    //             "failed on attempt {attempt}"
-    //         );
-    //         if should_validate {
-    //             // Now that we've triggered a device check, we can give up
-    //             break;
-    //         }
-    //     }
-    // }
+    /// Test whether the relay selector selects wireguard often enough, given no special constraints.
+    /// A Wireguard relay must be used to verify that the device is valid.
+    #[test]
+    fn test_validates_by_default() {
+        use mullvad_relay_selector::{RelaySelector, SelectorConfig};
+        use talpid_types::net::TunnelType;
+        // No special relay constraints / user settings are assumed
+        let config = SelectorConfig::default();
+        // It ought to be enough that the first 3 connection attempts will yield a Wireguard relay
+        for attempt in 0..3 {
+            let typ = RelaySelector::would_return(attempt, &config).unwrap();
+            assert_eq!(typ, TunnelType::Wireguard);
+        }
+    }
 }
diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index f8583ab4adfa..4c9a78025af4 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -433,6 +433,30 @@ impl RelaySelector {
         self.get_relay_with_order(&RETRY_ORDER, retry_attempt)
     }
 
+    /// Peek at which [`TunnelType`] that would be returned for a certain connection attempt for a given
+    /// [`SelectorConfig`]. Returns [`Option::None`] if the given config would return a custom
+    /// tunnel endpoint.
+    ///
+    /// # Note
+    /// This function is only really useful for testing-purposes. It is exposed to ease testing of
+    /// other mullvad crates which depend on the retry behaviour of [`RelaySelector`].
+    pub fn would_return(connection_attempt: usize, config: &SelectorConfig) -> Option<TunnelType> {
+        let config = SpecializedSelectorConfig::from(config);
+        match config {
+            // This case is not really interesting
+            SpecializedSelectorConfig::Custom(_) => None,
+            SpecializedSelectorConfig::Normal(config) => Some(
+                Self::pick_and_merge_query(
+                    &RETRY_ORDER,
+                    RelayQuery::from(config),
+                    connection_attempt,
+                )
+                .tunnel_protocol
+                .unwrap_or(TunnelType::Wireguard),
+            ),
+        }
+    }
+
     /// Returns a random relay and relay endpoint matching the current constraints defined by
     /// `retry_order` corresponsing to `retry_attempt`.
     pub fn get_relay_with_order(