Add location overrides per test to test-manager config

5 files changed, 104 insertions, 204 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index ff581c8f62..78a7cfacdf 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -17,6 +17,7 @@ use mullvad_relay_selector::{
 };
 use mullvad_types::{
     constraints::Constraint,
+    custom_list::CustomList,
     relay_constraints::{
         GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
     },
@@ -1199,160 +1200,48 @@ fn parse_am_i_mullvad(result: String) -> anyhow::Result<bool> {
     })
 }
 
-pub mod custom_lists {
-    use super::*;
-
-    use mullvad_types::custom_list::{CustomList, Id};
-    use std::sync::{LazyLock, Mutex};
-
-    // Expose all custom list variants as a shorthand.
-    pub use List::*;
-
-    /// The default custom list to use as location for all tests.
-    pub const DEFAULT_LIST: List = List::Nordic;
-
-    /// Mapping between [List] to daemon custom lists. Since custom list ids are assigned by the
-    /// daemon at the creation of the custom list settings object, we can't map a custom list
-    /// name to a specific list before runtime.
-    static IDS: LazyLock<Mutex<HashMap<List, Id>>> = LazyLock::new(|| Mutex::new(HashMap::new()));
-
-    /// Pre-defined (well-typed) custom lists which may be useful in different test scenarios.
-    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
-    pub enum List {
-        /// A selection of Nordic servers
-        Nordic,
-        /// A selection of European servers
-        Europe,
-        /// This custom list contains relays which are close geographically to the computer running
-        /// the test scenarios, which hopefully means there will be little latency between the test
-        /// machine and these relays
-        LowLatency,
-        /// Antithesis of [List::LowLatency], these relays are located far away from the test
-        /// server. Use this custom list if you want to simulate scenarios where the probability
-        /// of experiencing high latencies is desirable.
-        HighLatency,
-    }
-
-    impl List {
-        pub fn name(self) -> String {
-            use List::*;
-            match self {
-                Nordic => "Nordic".to_string(),
-                Europe => "Europe".to_string(),
-                LowLatency => "Low Latency".to_string(),
-                HighLatency => "High Latency".to_string(),
-            }
-        }
-
-        /// Iterator over all custom lists.
-        pub fn all() -> impl Iterator<Item = List> {
-            use List::*;
-            [Nordic, Europe, LowLatency, HighLatency].into_iter()
-        }
-
-        pub fn locations(self) -> impl Iterator<Item = GeographicLocationConstraint> {
-            use List::*;
-            let country = GeographicLocationConstraint::country;
-            let city = GeographicLocationConstraint::city;
-            match self {
-                Nordic => {
-                    vec![country("no"), country("se"), country("fi"), country("dk")].into_iter()
-                }
-                Europe => vec![
-                    // North
-                    country("se"),
-                    // West
-                    country("fr"),
-                    // East
-                    country("ro"),
-                    // South
-                    country("it"),
-                ]
-                .into_iter(),
-                LowLatency => {
-                    // Assumption: Test server is located in Gothenburg, Sweden.
-                    vec![city("se", "got")].into_iter()
-                }
-                HighLatency => {
-                    // Assumption: Test server is located in Gothenburg, Sweden.
-                    vec![country("au"), country("ca"), country("za")].into_iter()
-                }
-            }
-        }
-
-        pub fn to_constraint(self) -> Option<LocationConstraint> {
-            let ids = IDS.lock().unwrap();
-            let id = ids.get(&self)?;
-            Some(LocationConstraint::CustomList { list_id: *id })
-        }
-    }
-
-    impl From<List> for LocationConstraint {
-        fn from(custom_list: List) -> Self {
-            // TODO: Is this _too_ unsound ??
-            custom_list.to_constraint().unwrap()
-        }
-    }
-
-    /// Add a set of custom lists which can be used in different test scenarios.
-    ///
-    /// See [`List`] for available custom lists.
-    pub async fn add_default_lists(mullvad_client: &mut MullvadProxyClient) -> anyhow::Result<()> {
-        for custom_list in List::all() {
-            let id = mullvad_client
-                .create_custom_list(custom_list.name())
-                .await?;
-            let mut daemon_dito = find_custom_list(mullvad_client, &custom_list.name()).await?;
-            assert_eq!(id, daemon_dito.id);
-            for locations in custom_list.locations() {
-                daemon_dito.locations.insert(locations);
-            }
-            mullvad_client.update_custom_list(daemon_dito).await?;
-            // Associate this custom list variant with a specific, runtime custom list id.
-            IDS.lock().unwrap().insert(custom_list, id);
-        }
-        Ok(())
-    }
-
-    /// Set the default location to the custom list specified by `DEFAULT_LIST`. This also includes
-    /// entry location for multihop. It does not, however, affect bridge location for OpenVPN.
-    /// This is for simplify, as bridges default to using the server closest to the exit anyway, and
-    /// OpenVPN is slated for removal.
-    pub async fn set_default_location(
-        mullvad_client: &mut MullvadProxyClient,
-    ) -> anyhow::Result<()> {
-        let constraints = get_custom_list_location_relay_constraints(DEFAULT_LIST);
+/// Set the location to the given [`LocationConstraint`]. This also includes
+/// entry location for multihop. It does not, however, affect bridge location for OpenVPN.
+/// This is for simplify, as bridges default to using the server closest to the exit anyway, and
+/// OpenVPN is slated for removal.
+///
+/// NOTE: Calling this from within a test will overwrite the default test lcoation specified in
+/// the settings.
+pub async fn set_location(
+    mullvad_client: &mut MullvadProxyClient,
+    location: impl Into<LocationConstraint>,
+) -> anyhow::Result<()> {
+    let constraints = get_location_relay_constraints(location.into());
 
-        mullvad_client
-            .set_relay_settings(constraints.into())
-            .await
-            .context("Failed to set relay settings")
-    }
+    mullvad_client
+        .set_relay_settings(constraints.into())
+        .await
+        .context("Failed to set relay settings")
+}
 
-    fn get_custom_list_location_relay_constraints(custom_list: List) -> RelayConstraints {
-        let wireguard_constraints = mullvad_types::relay_constraints::WireguardConstraints {
-            entry_location: Constraint::Only(custom_list.into()),
-            ..Default::default()
-        };
+fn get_location_relay_constraints(custom_list: LocationConstraint) -> RelayConstraints {
+    let wireguard_constraints = mullvad_types::relay_constraints::WireguardConstraints {
+        entry_location: Constraint::Only(custom_list.clone()),
+        ..Default::default()
+    };
 
-        RelayConstraints {
-            location: Constraint::Only(custom_list.into()),
-            wireguard_constraints,
-            ..Default::default()
-        }
+    RelayConstraints {
+        location: Constraint::Only(custom_list),
+        wireguard_constraints,
+        ..Default::default()
     }
+}
 
-    /// Dig out a custom list from the daemon settings based on the custom list's name.
-    /// There should be an rpc for this.
-    async fn find_custom_list(
-        rpc: &mut MullvadProxyClient,
-        name: &str,
-    ) -> anyhow::Result<CustomList> {
-        rpc.get_settings()
-            .await?
-            .custom_lists
-            .into_iter()
-            .find(|list| list.name == name)
-            .ok_or(anyhow!("List '{name}' not found"))
-    }
+/// Dig out a custom list from the daemon settings based on the custom list's name.
+/// There should be an rpc for this.
+pub async fn find_custom_list(
+    rpc: &mut MullvadProxyClient,
+    name: &str,
+) -> anyhow::Result<CustomList> {
+    rpc.get_settings()
+        .await?
+        .custom_lists
+        .into_iter()
+        .find(|list| list.name == name)
+        .ok_or(anyhow!("List '{name}' not found"))
 }
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 9eb23fe13e..b369b28067 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -16,6 +16,7 @@ mod tunnel_state;
 mod ui;
 
 use itertools::Itertools;
+use mullvad_types::relay_constraints::{GeographicLocationConstraint, LocationConstraint};
 pub use test_metadata::TestMetadata;
 
 use anyhow::Context;
@@ -27,7 +28,7 @@ use crate::{
     package::get_version_from_path,
 };
 use config::TEST_CONFIG;
-use helpers::{get_app_env, install_app};
+use helpers::{find_custom_list, get_app_env, install_app, set_location};
 pub use install::test_upgrade_app;
 use mullvad_management_interface::MullvadProxyClient;
 use test_rpc::{meta::Os, ServiceClient};
@@ -137,7 +138,7 @@ pub fn get_filtered_tests(specified_tests: &[String]) -> Result<Vec<TestMetadata
 pub async fn prepare_daemon(
     rpc: &ServiceClient,
     rpc_provider: &RpcClientProvider,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<MullvadProxyClient> {
     // Check if daemon should be restarted
     let mut mullvad_client = ensure_daemon_version(rpc, rpc_provider)
         .await
@@ -152,9 +153,60 @@ pub async fn prepare_daemon(
         .await
         .context("Failed to disconnect daemon after test")?;
     helpers::ensure_logged_in(&mut mullvad_client).await?;
-    helpers::custom_lists::add_default_lists(&mut mullvad_client).await?;
-    helpers::custom_lists::set_default_location(&mut mullvad_client).await?;
 
+    Ok(mullvad_client)
+}
+
+/// Create and selects an "anonymous" custom list for this test. The custom list will
+/// have the same name as the test and contain the locations as specified by
+/// [`TestMetadata`] location field.
+pub async fn set_test_location(
+    mullvad_client: &mut MullvadProxyClient,
+    test: &TestMetadata,
+) -> anyhow::Result<()> {
+    // If no location is specified for the test, don't do anything and use the default value of the app
+    let Some(locations) = test.location.as_ref() else {
+        return Ok(());
+    };
+    // Convert locations from the test config to actual location constraints
+    let locations: Vec<GeographicLocationConstraint> = locations
+        .iter()
+        .map(|input| {
+            input
+                .parse::<GeographicLocationConstraint>()
+                .with_context(|| format!("Failed to parse {input}"))
+        })
+        .try_collect()?;
+
+    log::debug!(
+        "Creating custom list {} with locations '{:?}'",
+        test.name,
+        locations
+    );
+
+    // Add the custom list to the current app instance
+    // NOTE: This const is actually defined in, `mullvad_types::custom_list`, but we cannot import it.
+    const CUSTOM_LIST_NAME_MAX_SIZE: usize = 30;
+    let mut custom_list_name = test.name.to_string();
+    custom_list_name.truncate(CUSTOM_LIST_NAME_MAX_SIZE);
+    log::debug!("Creating custom list {custom_list_name} with locations '{locations:?}'");
+
+    let list_id = mullvad_client
+        .create_custom_list(custom_list_name.clone())
+        .await?;
+
+    let mut custom_list = find_custom_list(mullvad_client, &custom_list_name).await?;
+
+    assert_eq!(list_id, custom_list.id);
+    for location in locations {
+        custom_list.locations.insert(location);
+    }
+    mullvad_client.update_custom_list(custom_list).await?;
+    log::debug!("Added custom list");
+
+    set_location(mullvad_client, LocationConstraint::CustomList { list_id })
+        .await
+        .with_context(|| format!("Failed to set location to custom list with ID '{list_id:?}'"))?;
     Ok(())
 }
 
diff --git a/test/test-manager/src/tests/test_metadata.rs b/test/test-manager/src/tests/test_metadata.rs
index 79c7f74def..961db4510a 100644
--- a/test/test-manager/src/tests/test_metadata.rs
+++ b/test/test-manager/src/tests/test_metadata.rs
@@ -1,13 +1,16 @@
 use super::TestWrapperFunction;
 use test_rpc::{meta::Os, mullvad_daemon::MullvadClientVersion};
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct TestMetadata {
     pub name: &'static str,
     pub targets: &'static [Os],
     pub mullvad_client_version: MullvadClientVersion,
     pub func: TestWrapperFunction,
+    /// Priority order of the tests, unless specific tests are given as the `TEST_FILTERS` argument
     pub priority: Option<i32>,
+    /// A list of location that will be used for by the test
+    pub location: Option<Vec<String>>,
 }
 
 // Register our test metadata struct with inventory to allow submitting tests of this type.
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index bc15ccb8a3..b64a16d854 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -187,15 +187,7 @@ pub async fn test_wireguard_over_shadowsocks(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
-    // This is an attempt to try to reduce this type of flakiness.
-    use helpers::custom_lists::LowLatency;
-
-    let query = RelayQueryBuilder::new()
-        .wireguard()
-        .shadowsocks()
-        .location(LowLatency)
-        .build();
+    let query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
@@ -294,16 +286,7 @@ pub async fn test_multihop(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
-    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
-    // This is an attempt to try to reduce this type of flakiness.
-    use helpers::custom_lists::LowLatency;
-
-    let query = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .location(LowLatency)
-        .entry(LowLatency)
-        .build();
+    let query = RelayQueryBuilder::new().wireguard().multihop().build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
@@ -455,10 +438,6 @@ pub async fn test_quantum_resistant_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
-    // This is an attempt to try to reduce this type of flakiness.
-    use helpers::custom_lists::LowLatency;
-
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::Off)
         .await
@@ -472,10 +451,7 @@ pub async fn test_quantum_resistant_tunnel(
 
     log::info!("Setting tunnel protocol to WireGuard");
 
-    let query = RelayQueryBuilder::new()
-        .wireguard()
-        .location(LowLatency)
-        .build();
+    let query = RelayQueryBuilder::new().wireguard().build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
@@ -536,10 +512,6 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
-    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
-    // This is an attempt to try to reduce this type of flakiness.
-    use helpers::custom_lists::LowLatency;
-
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::On)
         .await
@@ -549,8 +521,6 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
         .wireguard()
         .multihop()
         .udp2tcp()
-        .entry(LowLatency)
-        .location(LowLatency)
         .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
@@ -577,10 +547,6 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    // NOTE: We have experienced flakiness due to timeout issues if distant relays are selected.
-    // This is an attempt to try to reduce this type of flakiness.
-    use helpers::custom_lists::LowLatency;
-
     mullvad_client
         .set_quantum_resistant_tunnel(wireguard::QuantumResistantState::On)
         .await
@@ -590,8 +556,6 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
         .wireguard()
         .multihop()
         .shadowsocks()
-        .entry(LowLatency)
-        .location(LowLatency)
         .build();
 
     apply_settings_from_relay_query(&mut mullvad_client, query).await?;
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index 6b9ce5b3a9..9491308eb8 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -91,19 +91,11 @@ pub async fn test_ui_tunnel_settings(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    // NOTE: This test connects multiple times using various settings, some of which may cause a
-    // significant increase in connection time, e.g. multihop and OpenVPN. For this reason, it is
-    // preferable to only target low latency servers.
-    use helpers::custom_lists::LowLatency;
-
     // tunnel-state.spec precondition: a single WireGuard relay should be selected
     log::info!("Select WireGuard relay");
     let entry = helpers::constrain_to_relay(
         &mut mullvad_client,
-        RelayQueryBuilder::new()
-            .wireguard()
-            .location(LowLatency)
-            .build(),
+        RelayQueryBuilder::new().wireguard().build(),
     )
     .await?;