Add location overrides per test to test-manager config

13 files changed, 305 insertions, 225 deletions

diff --git a/test/Cargo.lock b/test/Cargo.lock
index a4e2ee9df7..064ae149ad 100644
--- a/test/Cargo.lock
+++ b/test/Cargo.lock
@@ -3597,6 +3597,7 @@ dependencies = [
  "dirs",
  "env_logger",
  "futures",
+ "glob",
  "hyper-util",
  "inventory",
  "ipnetwork",
diff --git a/test/test-manager/Cargo.toml b/test/test-manager/Cargo.toml
index 2671ea454a..3310ab770f 100644
--- a/test/test-manager/Cargo.toml
+++ b/test/test-manager/Cargo.toml
@@ -32,6 +32,7 @@ async-trait = { workspace = true }
 uuid = "1.3"
 dirs = "5.0.1"
 scopeguard = "1.2"
+glob = "0.3"
 
 serde = { workspace = true }
 serde_json = { workspace = true }
diff --git a/test/test-manager/src/config/manifest.rs b/test/test-manager/src/config/manifest.rs
index 8ef96dc89d..f0ff1419a9 100644
--- a/test/test-manager/src/config/manifest.rs
+++ b/test/test-manager/src/config/manifest.rs
@@ -1,28 +1,35 @@
-//! Config definition.
-//! TODO: Document struct and link to that documentation
+//! Config definition, see [`Config`].
 
 use std::collections::BTreeMap;
 
 use serde::{Deserialize, Serialize};
+use test_locations::TestLocationList;
 
 use super::VmConfig;
 use crate::tests::config::DEFAULT_MULLVAD_HOST;
 
-#[derive(Default, Serialize, Deserialize, Clone)]
+/// Global configuration for the `test-manager`.
+///
+/// Can be modified using either the setting file, see
+/// [`crate::config::io::ConfigFile::get_config_path`] or
+/// the `test-manager config` CLI subcommand.
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub struct Config {
     #[serde(skip)]
     pub runtime_opts: RuntimeOptions,
     pub vms: BTreeMap<String, VmConfig>,
     pub mullvad_host: Option<String>,
+    #[serde(default)]
+    pub test_locations: TestLocationList,
 }
 
-#[derive(Default, Serialize, Deserialize, Clone)]
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub struct RuntimeOptions {
     pub display: Display,
     pub keep_changes: bool,
 }
 
-#[derive(Default, Serialize, Deserialize, Clone)]
+#[derive(Debug, Default, Serialize, Deserialize, Clone)]
 pub enum Display {
     #[default]
     None,
@@ -45,3 +52,167 @@ impl Config {
         })
     }
 }
+
+mod test_locations {
+    use serde::{
+        de::{Deserialize, Deserializer, Error, MapAccess, Visitor},
+        ser::{Serialize, SerializeMap},
+        Deserialize as DeserDerive, Serialize as SerDerive,
+    };
+    use std::fmt;
+
+    #[derive(Clone, Default)]
+    pub struct TestLocation(glob::Pattern, Vec<String>);
+
+    impl fmt::Debug for TestLocation {
+        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+            write!(f, "{}: {:?}", self.0, &self.1)
+        }
+    }
+
+    /// Relay/location overrides for tests.
+    ///
+    /// # Deserializing with `serde-json`
+    ///
+    /// The format is a list of maps with a single key-value
+    /// pair, where the key is a glob pattern that will be matched against the test name, and the
+    /// value is a list of locations to use for that test. The first match will be used.
+    ///
+    /// Example:
+    /// ```json
+    /// {
+    ///   // other fields
+    ///    "test_locations": [
+    ///        { "*daita*": [ "se-got-wg-001", "se-got-wg-002" ] },
+    ///        { "*": [ "se" ] }
+    ///    ]
+    /// }
+    /// ```
+    ///
+    /// The above example will set the locations for the test `test_daita` to a custom list
+    /// containing `se-got-wg-001` and `se-got-wg-002`. The `*` is a wildcard that will match
+    /// any test name. The order of the list is important, as the first match will be used.
+    #[derive(Debug, DeserDerive, SerDerive, Clone, Default)]
+    pub struct TestLocationList(pub Vec<TestLocation>);
+
+    impl TestLocationList {
+        pub fn lookup(&self, test: &str) -> Option<&Vec<String>> {
+            self.0
+                .iter()
+                .find(|TestLocation(test_glob, _)| test_glob.matches(test))
+                .map(|TestLocation(_, locations)| locations)
+        }
+    }
+
+    struct TestLocationVisitor;
+
+    impl<'de> Visitor<'de> for TestLocationVisitor {
+        // The type that our Visitor is going to produce.
+        type Value = TestLocation;
+
+        // Format a message stating what data this Visitor expects to receive.
+        fn expecting(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
+            formatter.write_str("A list of maps")
+        }
+
+        fn visit_map<M>(self, mut access: M) -> Result<Self::Value, M::Error>
+        where
+            M: MapAccess<'de>,
+        {
+            let (key, value) =
+                access
+                    .next_entry::<String, Vec<String>>()?
+                    .ok_or(M::Error::custom(
+                    "Test location map should contain exactly one key-value pair, but it was empty",
+                ))?;
+            let glob = glob::Pattern::new(&key).map_err(|err| {
+                M::Error::custom(format!(
+                    "Cannot compile glob pattern from: {key} error: {err:?}"
+                ))
+            })?;
+
+            if let Some((key, value)) = access.next_entry::<String, Vec<String>>()? {
+                return Err(M::Error::custom(format!(
+                    "Test location map should contain exactly one key-value pair, but found another key: '{key}' and value: '{value:?}'"
+                )));
+            }
+
+            Ok(TestLocation(glob, value))
+        }
+    }
+
+    impl<'de> Deserialize<'de> for TestLocation {
+        fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+        where
+            D: Deserializer<'de>,
+        {
+            deserializer.deserialize_map(TestLocationVisitor)
+        }
+    }
+
+    impl Serialize for TestLocation {
+        fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+        where
+            S: serde::Serializer,
+        {
+            let mut map = serializer.serialize_map(Some(1))?;
+            map.serialize_entry(self.0.as_str(), &self.1)?;
+            map.end()
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parse_test_location_empty() {
+        let config = r#"
+            {
+                "vms": {},
+                "mullvad_host": "mullvad.net"
+            }"#;
+
+        let config: Config = serde_json::from_str(config).unwrap();
+        assert!(config.test_locations.0.is_empty());
+    }
+
+    #[test]
+    fn parse_test_location_not_empty() {
+        let config = r#"
+            {
+                "vms": {},
+                "mullvad_host": "mullvad.net",
+                "test_locations": [
+                    { "*daita": [ "se-got-wg-001", "se-got-wg-002" ] },
+                    { "*": [ "se" ] }
+                ]
+            }"#;
+
+        let config: Config = serde_json::from_str(config).unwrap();
+        assert!(config
+            .test_locations
+            .lookup("test_daita")
+            .unwrap()
+            .contains(&"se-got-wg-002".to_string()));
+        assert!(!config.test_locations.0.is_empty());
+    }
+
+    #[test]
+    fn parse_multiple_keys_in_map_should_fail() {
+        let config = r#"
+            {
+                "vms": {},
+                "mullvad_host": "mullvad.net",
+                "test_locations": [
+                    {
+                        "*daita": [ "se-got-wg-001", "se-got-wg-002" ],
+                        "*test": ["se"]
+                    },
+                ]
+            }"#;
+
+        let _err = serde_json::from_str::<Config>(config).unwrap_err();
+    }
+}
diff --git a/test/test-manager/src/config/mod.rs b/test/test-manager/src/config/mod.rs
index 4dc6e72482..4b510b87cc 100644
--- a/test/test-manager/src/config/mod.rs
+++ b/test/test-manager/src/config/mod.rs
@@ -9,11 +9,3 @@ use error::Error;
 pub use io::ConfigFile;
 pub use manifest::{Config, Display};
 pub use vm::{Architecture, OsType, PackageType, Provisioner, VmConfig, VmType};
-
-#[cfg(test)]
-mod tests {
-    use super::*;
-
-    #[test]
-    fn parse_relay_location_per_test_override() {}
-}
diff --git a/test/test-manager/src/main.rs b/test/test-manager/src/main.rs
index 95274e9a7a..cdee85e44a 100644
--- a/test/test-manager/src/main.rs
+++ b/test/test-manager/src/main.rs
@@ -266,7 +266,12 @@ async fn main() -> Result<()> {
             };
 
             if let Some(mullvad_host) = mullvad_host {
-                log::trace!("Setting Mullvad host using --mullvad-host flag");
+                match config.mullvad_host {
+                    Some(old_host) => {
+                        log::info!("Overriding Mullvad host from {old_host} to {mullvad_host}",)
+                    }
+                    None => log::info!("Setting Mullvad host to {mullvad_host}",),
+                };
                 config.mullvad_host = Some(mullvad_host);
             }
             let mullvad_host = config.get_host();
@@ -327,7 +332,11 @@ async fn main() -> Result<()> {
                 test_rpc::meta::Os::from(vm_config.os_type),
                 openvpn_certificate,
             ));
-            let tests = get_filtered_tests(&test_filters)?;
+
+            let mut tests = get_filtered_tests(&test_filters)?;
+            for test in tests.iter_mut() {
+                test.location = config.test_locations.lookup(test.name).cloned();
+            }
 
             // For convenience, spawn a SOCKS5 server that is reachable for tests that need it
             let socks = socks_server::spawn(SocketAddr::new(
diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 88577fb3b3..3d2f1702e0 100644
--- a/test/test-manager/src/run_tests.rs
+++ b/test/test-manager/src/run_tests.rs
@@ -157,15 +157,20 @@ pub async fn run(
     };
 
     for test in tests {
-        tests::prepare_daemon(&test_runner_client, &rpc_provider)
+        let mut mullvad_client = tests::prepare_daemon(&test_runner_client, &rpc_provider)
             .await
             .context("Failed to reset daemon before test")?;
 
-        let mullvad_client = rpc_provider
-            .mullvad_client(test.mullvad_client_version)
-            .await;
+        tests::set_test_location(&mut mullvad_client, &test)
+            .await
+            .context("Failed to create custom list from test locations")?;
+
         test_handler
-            .run_test(&test.func, test.name, mullvad_client)
+            .run_test(
+                &test.func,
+                test.name,
+                MullvadClientArgument::WithClient(mullvad_client),
+            )
             .await?;
     }
 
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?;
 
diff --git a/test/test-manager/test_macro/src/lib.rs b/test/test-manager/test_macro/src/lib.rs
index cddb6c5a2f..e16968bd9d 100644
--- a/test/test-manager/test_macro/src/lib.rs
+++ b/test/test-manager/test_macro/src/lib.rs
@@ -191,6 +191,7 @@ fn create_test(test_function: TestFunction) -> proc_macro2::TokenStream {
             mullvad_client_version: #function_mullvad_version,
             func: #wrapper_closure,
             priority: #test_function_priority,
+            location: None,
         });
     }
 }
diff --git a/test/test-rpc/src/mullvad_daemon.rs b/test/test-rpc/src/mullvad_daemon.rs
index 10cc00c3fc..e9865440b9 100644
--- a/test/test-rpc/src/mullvad_daemon.rs
+++ b/test/test-rpc/src/mullvad_daemon.rs
@@ -27,7 +27,7 @@ pub enum Verbosity {
     Trace,
 }
 
-#[derive(Clone, Copy, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum MullvadClientVersion {
     None,
     New,