Merge branch 'fix-overwrite-relay-location-test-framework'

6 files changed, 256 insertions, 171 deletions

diff --git a/test/test-manager/src/tests/daita.rs b/test/test-manager/src/tests/daita.rs
index 24f5f4b0df..91687002aa 100644
--- a/test/test-manager/src/tests/daita.rs
+++ b/test/test-manager/src/tests/daita.rs
@@ -3,8 +3,8 @@ use futures::StreamExt;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
-    relay_constraints::GeographicLocationConstraint, relay_list::RelayEndpointData,
-    states::TunnelState,
+    constraints::Constraint, relay_constraints::GeographicLocationConstraint,
+    relay_list::RelayEndpointData, states::TunnelState,
 };
 use talpid_types::{net::TunnelEndpoint, tunnel::ErrorStateCause};
 use test_macro::test_function;
@@ -29,13 +29,11 @@ pub async fn test_daita(
     _rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    let relay_list = mullvad_client.get_relay_locations().await?;
-    let wg_relays = relay_list
-        .relays()
-        .flat_map(|relay| match &relay.endpoint_data {
-            RelayEndpointData::Wireguard(wireguard) => Some((relay, wireguard)),
-            _ => None,
-        });
+    let relays = helpers::get_all_pickable_relays(&mut mullvad_client).await?;
+    let wg_relays = relays.iter().flat_map(|relay| match &relay.endpoint_data {
+        RelayEndpointData::Wireguard(wireguard) => Some((relay, wireguard)),
+        _ => None,
+    });
 
     // Select two relays to use for the test, one with DAITA and one without.
     let daita_relay = wg_relays
@@ -62,38 +60,24 @@ pub async fn test_daita(
     );
     log::info!("Selected non-daita relay: {}", non_daita_relay.hostname);
 
-    let non_daita_location_query = RelayQueryBuilder::new()
-        .wireguard()
-        .location(non_daita_relay_location.clone())
-        .build();
-
-    let daita_location_query = RelayQueryBuilder::new()
-        .wireguard()
-        .location(daita_relay_location.clone())
-        .build();
-
-    let daita_to_non_daita_multihop_query = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .entry(daita_relay_location.clone())
-        .location(non_daita_relay_location.clone())
-        .build();
-
-    let non_daita_multihop_query = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .entry(non_daita_relay_location.clone())
-        .build();
+    log::info!("Setting wireguard and DAITA");
+    let wireguard_query = RelayQueryBuilder::new().wireguard().build();
+    helpers::apply_settings_from_relay_query(&mut mullvad_client, wireguard_query.clone()).await?;
+    mullvad_client.set_enable_daita(true).await?;
 
     let mut events = mullvad_client
         .events_listen()
         .await?
         .inspect(|event| log::debug!("New daemon event: {event:?}"));
 
-    log::info!("Connecting to non-daita relay with DAITA by automatically using multihop");
+    log::info!("Connecting to non-daita relay with DAITA should automatically use multihop");
     {
-        helpers::set_relay_settings(&mut mullvad_client, non_daita_location_query.clone()).await?;
-        mullvad_client.set_enable_daita(true).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(false).await?;
+
         mullvad_client.connect_tunnel().await?;
         let state = wait_for_daemon_reconnect(&mut events)
             .await
@@ -106,8 +90,12 @@ pub async fn test_daita(
         log::info!("Successfully multihopped with 'direct only' disabled");
     }
 
-    log::info!("Connecting to non-daita relay with 'DAITA: direct only'");
+    log::info!("Connecting to non-daita relay with 'direct_only' shoud fail");
     {
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+        })
+        .await?;
         mullvad_client.set_daita_direct_only(true).await?;
 
         let result = wait_for_daemon_reconnect(&mut events).await;
@@ -121,9 +109,13 @@ pub async fn test_daita(
         log::info!("Failed to connect, this is expected!");
     }
 
-    log::info!("Connecting to daita relay with 'direct_only' disabled");
+    log::info!("Connecting to daita relay with 'direct_only' should not use multihop");
     {
-        helpers::set_relay_settings(&mut mullvad_client, daita_location_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(daita_relay_location.clone().into());
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
 
         let state = wait_for_daemon_reconnect(&mut events)
             .await
@@ -139,9 +131,17 @@ pub async fn test_daita(
         log::info!("Successfully singlehopped with 'direct_only' disabled");
     }
 
-    log::info!("Connecting to daita relay with multihop");
+    log::info!("Connecting to a daita relay as entry for multihop and `direct_only` should work");
     {
-        helpers::set_relay_settings(&mut mullvad_client, daita_to_non_daita_multihop_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(non_daita_relay_location.clone().into());
+            constraint.wireguard_constraints.entry_location =
+                Constraint::Only(daita_relay_location.clone().into());
+            constraint.wireguard_constraints.use_multihop = true;
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
+
         let state = wait_for_daemon_reconnect(&mut events)
             .await
             .context("Failed to connect via daita location with multihop enabled")?;
@@ -153,9 +153,19 @@ pub async fn test_daita(
         log::info!("Successfully connected with multihop");
     }
 
-    log::info!("Connecting to non_daita relay with multihop");
+    log::info!(
+        "Connecting to a non daita relay as entry for multihop and `direct_only` should fail"
+    );
     {
-        helpers::set_relay_settings(&mut mullvad_client, non_daita_multihop_query).await?;
+        helpers::update_relay_constraints(&mut mullvad_client, |constraint| {
+            constraint.location = Constraint::Only(daita_relay_location.clone().into());
+            constraint.wireguard_constraints.entry_location =
+                Constraint::Only(non_daita_relay_location.into());
+            constraint.wireguard_constraints.use_multihop = true;
+        })
+        .await?;
+        mullvad_client.set_daita_direct_only(true).await?;
+
         let result = wait_for_daemon_reconnect(&mut events).await;
         let Err(Error::UnexpectedErrorState(state)) = result else {
             bail!("Connection failed unsuccessfully, reason: {:?}", result);
diff --git a/test/test-manager/src/tests/dns.rs b/test/test-manager/src/tests/dns.rs
index 28db345cf7..b54be0e02e 100644
--- a/test/test-manager/src/tests/dns.rs
+++ b/test/test-manager/src/tests/dns.rs
@@ -8,7 +8,6 @@ use std::{
 use itertools::Itertools;
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_types::{
-    relay_constraints::RelaySettings,
     settings,
     wireguard::{DaitaSettings, QuantumResistantState},
     ConnectionConfig, CustomTunnelEndpoint,
@@ -18,7 +17,7 @@ use test_macro::test_function;
 use test_rpc::ServiceClient;
 
 use super::{
-    helpers::{self, connect_and_wait, set_relay_settings},
+    helpers::{self, connect_and_wait, set_custom_endpoint},
     Error, TestContext,
 };
 use crate::{
@@ -658,7 +657,7 @@ async fn connect_local_wg_relay(mullvad_client: &mut MullvadProxyClient) -> Resu
         CUSTOM_TUN_REMOTE_REAL_PORT,
     );
 
-    let relay_settings = RelaySettings::CustomTunnelEndpoint(CustomTunnelEndpoint {
+    let custom_tunnel_endpoint = CustomTunnelEndpoint {
         host: peer_addr.ip().to_string(),
         config: ConnectionConfig::Wireguard(wireguard::ConnectionConfig {
             tunnel: wireguard::TunnelConfig {
@@ -678,9 +677,8 @@ async fn connect_local_wg_relay(mullvad_client: &mut MullvadProxyClient) -> Resu
             fwmark: None,
             ipv6_gateway: None,
         }),
-    });
-
-    set_relay_settings(mullvad_client, relay_settings)
+    };
+    set_custom_endpoint(mullvad_client, custom_tunnel_endpoint)
         .await
         .expect("failed to update relay settings");
 
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 195c3da26a..f4a5986755 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -13,7 +13,8 @@ use anyhow::{anyhow, bail, ensure, Context};
 use futures::StreamExt;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_relay_selector::{
-    query::RelayQuery, GetRelay, RelaySelector, SelectorConfig, WireguardConfig,
+    query::{OpenVpnRelayQuery, RelayQuery, WireguardRelayQuery},
+    GetRelay, RelaySelector, SelectorConfig, WireguardConfig,
 };
 use mullvad_types::{
     constraints::Constraint,
@@ -133,7 +134,7 @@ pub async fn reboot(rpc: &mut ServiceClient) -> Result<(), Error> {
     #[cfg(target_os = "macos")]
     crate::vm::network::macos::configure_tunnel()
         .await
-        .map_err(|error| Error::Other(format!("Failed to recreate custom wg tun: {error}")))?;
+        .context("Failed to recreate custom wg tun: {error}")?;
 
     Ok(())
 }
@@ -596,40 +597,70 @@ impl<T> Drop for AbortOnDrop<T> {
     }
 }
 
+/// Applies the given query to the daemon location selection. The query will be intersected with
+/// the current location settings, so that the default location custom list is still used.
 pub async fn apply_settings_from_relay_query(
     mullvad_client: &mut MullvadProxyClient,
     query: RelayQuery,
-) -> Result<(), Error> {
-    let (constraints, bridge_state, bridge_settings, obfuscation) = query.into_settings();
+) -> anyhow::Result<()> {
+    // To prevent overwriting default custom list location constraint, we make an intersection with
+    // a query containing only the current location constraint
+    let intersected_relay_query = intersect_with_current_location(mullvad_client, query.clone())
+        .await
+        .with_context(|| {
+            format!("Failed to join query with current daemon settings. Query: {query:#?}")
+        })?;
+    let (constraints, bridge_state, bridge_settings, obfuscation) =
+        intersected_relay_query.into_settings();
 
     mullvad_client
         .set_relay_settings(constraints.into())
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))?;
+        .context("Failed to set daemon settings")?;
     mullvad_client
         .set_bridge_state(bridge_state)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set bridge state: {}", error)))?;
+        .context("Failed to set bridge state")?;
     mullvad_client
         .set_bridge_settings(bridge_settings)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set bridge settings: {}", error)))?;
+        .context("Failed to set bridge settings")?;
     mullvad_client
         .set_obfuscation_settings(obfuscation)
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set obfuscation settings: {}", error)))
+        .context("Failed to set obfuscation settings")?;
+    Ok(())
 }
 
-pub async fn set_relay_settings(
+pub async fn set_custom_endpoint(
     mullvad_client: &mut MullvadProxyClient,
-    relay_settings: impl Into<RelaySettings>,
+    custom_endpoint: mullvad_types::CustomTunnelEndpoint,
 ) -> Result<(), Error> {
     mullvad_client
-        .set_relay_settings(relay_settings.into())
+        .set_relay_settings(RelaySettings::CustomTunnelEndpoint(custom_endpoint))
         .await
         .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))
 }
 
+pub async fn update_relay_constraints(
+    mullvad_client: &mut MullvadProxyClient,
+    fn_mut: impl FnOnce(&mut RelayConstraints),
+) -> anyhow::Result<()> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get setting from daemon")?;
+    let RelaySettings::Normal(mut relay_constraints) = settings.relay_settings else {
+        bail!("Mutating custom endpoint not supported");
+    };
+    fn_mut(&mut relay_constraints);
+    mullvad_client
+        .set_relay_settings(RelaySettings::Normal(relay_constraints))
+        .await
+        .context("Failed to set relay settings")?;
+    Ok(())
+}
+
 /// Wait for the relay list to be updated, to make sure we have the overridden one.
 /// Time out after a while.
 pub async fn ensure_updated_relay_list(
@@ -695,18 +726,99 @@ pub async fn get_app_env() -> anyhow::Result<HashMap<String, String>> {
     ]))
 }
 
-/// Constrain the daemon to only select the relay selected with `query` when establishing all
-/// future tunnels (until relay settings are updated, see [`set_relay_settings`]). Returns the
-/// selected [`Relay`] for future reference.
+/// Constrain the daemon to only select the relay compatible with `query` and the current relay
+/// settings when establishing all future tunnels (until relay settings are updated, see [`set_relay_settings`]).
+/// Returns the selected [`Relay`] for future reference.
+pub async fn constrain_to_relay(
+    mullvad_client: &mut MullvadProxyClient,
+    query: RelayQuery,
+) -> anyhow::Result<Relay> {
+    let intersect_query = intersect_with_current_location(mullvad_client, query).await?;
+    let (exit, relay_constraints) =
+        get_single_relay_location_contraint(mullvad_client, intersect_query).await?;
+
+    update_relay_constraints(mullvad_client, |current_constraints| {
+        *current_constraints = relay_constraints
+    })
+    .await
+    .unwrap();
+
+    Ok(exit)
+}
+
+/// Intersects the given query with the current location constraints, to prevent accidentally
+/// overwriting the default location custom list
+async fn intersect_with_current_location(
+    mullvad_client: &mut MullvadProxyClient,
+    query: RelayQuery,
+) -> anyhow::Result<RelayQuery> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get settings")?;
+    let RelaySettings::Normal(constraint) = settings.relay_settings else {
+        unimplemented!("Setting location for a custom endpoint is not supported");
+    };
+
+    // Construct a relay query preserving only the information about the current location
+    let current_location_query = RelayQuery::new(
+        constraint.location,
+        Constraint::Any,
+        Constraint::Any,
+        Constraint::Any,
+        WireguardRelayQuery {
+            entry_location: constraint.wireguard_constraints.entry_location,
+            ..Default::default()
+        },
+        OpenVpnRelayQuery {
+            bridge_settings: mullvad_relay_selector::query::BridgeQuery::Normal(
+                settings.bridge_settings.normal,
+            ),
+            ..Default::default()
+        },
+    )?;
+    use mullvad_types::Intersection;
+    let intersect_query = query
+        .intersection(current_location_query)
+        .context("Relay query incompatible with default settings")?;
+    Ok(intersect_query)
+}
+
+/// Get a query representing the current daemon settings
+async fn get_query_from_current_settings(
+    mullvad_client: &mut MullvadProxyClient,
+) -> anyhow::Result<RelayQuery> {
+    let settings = mullvad_client
+        .get_settings()
+        .await
+        .context("Failed to get settings")?;
+    RelayQuery::try_from(settings).context("Failed to convert settings to relay query")
+}
+
+pub async fn get_all_pickable_relays(
+    mullvad_client: &mut MullvadProxyClient,
+) -> anyhow::Result<Vec<Relay>> {
+    let settings = mullvad_client.get_settings().await?;
+    let relay_list = mullvad_client.get_relay_locations().await?;
+    let relays = mullvad_relay_selector::filter_matching_relay_list(
+        &helpers::get_query_from_current_settings(mullvad_client).await?,
+        &relay_list,
+        &settings.custom_lists,
+    );
+    Ok(relays)
+}
+
+/// Selects a relay compatible with the given query and relay list from the client, and returns a
+/// location constraint for only that relay, along with the relay itself.
 ///
 /// # Note
 /// This function does not handle bridges and multihop configurations (currently). There is no
 /// particular reason for this other than it not being needed at the time, so feel free to extend
 /// this function :).
-pub async fn constrain_to_relay(
+async fn get_single_relay_location_contraint(
     mullvad_client: &mut MullvadProxyClient,
     query: RelayQuery,
-) -> anyhow::Result<Relay> {
+) -> anyhow::Result<(Relay, RelayConstraints)> {
     /// Convert the result of invoking the relay selector to a relay constraint.
     fn convert_to_relay_constraints(
         query: RelayQuery,
@@ -726,18 +838,11 @@ pub async fn constrain_to_relay(
             unsupported => bail!("Can not constrain to a {unsupported:?}"),
         }
     }
-
     let settings = mullvad_client.get_settings().await?;
-    // Construct a relay selector with up-to-date information from the runnin daemon's relay list
     let relay_list = mullvad_client.get_relay_locations().await?;
     let relay_selector = get_daemon_relay_selector(&settings, relay_list);
-    // Select an(y) appropriate relay for the given query and constrain the daemon to only connect
-    // to that specific relay (when connecting).
     let relay = relay_selector.get_relay_by_query(query.clone())?;
-    let (exit, relay_constraints) = convert_to_relay_constraints(query, relay)?;
-    set_relay_settings(mullvad_client, RelaySettings::Normal(relay_constraints)).await?;
-
-    Ok(exit)
+    convert_to_relay_constraints(query, relay)
 }
 
 /// Get a mirror of the relay selector used by the daemon.
@@ -1200,36 +1305,32 @@ fn parse_am_i_mullvad(result: String) -> anyhow::Result<bool> {
     })
 }
 
-/// 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.
+/// Set the location to the given [`LocationConstraint`]. The same location constraint will be set
+/// for the multihop entry and OpenVPN bridge location as well.
 pub async fn set_location(
     mullvad_client: &mut MullvadProxyClient,
     location: impl Into<LocationConstraint>,
 ) -> anyhow::Result<()> {
-    let constraints = get_location_relay_constraints(location.into());
+    let location_constraint: LocationConstraint = location.into();
+    let mut settings = mullvad_client
+        .get_settings()
+        .await
+        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))?;
 
+    settings.bridge_settings.normal.location = Constraint::Only(location_constraint.clone());
     mullvad_client
-        .set_relay_settings(constraints.into())
-        .await
-        .context("Failed to set relay settings")
-}
+        .set_bridge_settings(settings.bridge_settings)
+        .await?;
 
-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()
+    let RelaySettings::Normal(mut constraint) = settings.relay_settings else {
+        unimplemented!("Setting location for a custom endpoint is not supported");
     };
-
-    RelayConstraints {
-        location: Constraint::Only(custom_list),
-        wireguard_constraints,
-        ..Default::default()
-    }
+    constraint.location = Constraint::Only(location_constraint.clone());
+    constraint.wireguard_constraints.entry_location = Constraint::Only(location_constraint);
+    mullvad_client
+        .set_relay_settings(RelaySettings::Normal(constraint))
+        .await?;
+    Ok(())
 }
 
 /// Dig out a custom list from the daemon settings based on the custom list's name.
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 4f5362167f..66635a10c1 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -72,9 +72,8 @@ pub enum Error {
     #[error("GUI test binary missing")]
     MissingGuiTest,
 
-    #[cfg(target_os = "macos")]
     #[error("An error occurred: {0}")]
-    Other(String),
+    Other(#[from] anyhow::Error),
 }
 
 #[derive(Clone)]
@@ -145,13 +144,13 @@ pub async fn prepare_daemon(
         .context("Failed to restart daemon")?;
 
     log::debug!("Resetting daemon settings before test");
+    helpers::disconnect_and_wait(&mut mullvad_client)
+        .await
+        .context("Failed to disconnect daemon after test")?;
     mullvad_client
         .reset_settings()
         .await
         .context("Failed to reset settings")?;
-    helpers::disconnect_and_wait(&mut mullvad_client)
-        .await
-        .context("Failed to disconnect daemon after test")?;
     helpers::ensure_logged_in(&mut mullvad_client).await?;
 
     Ok(mullvad_client)
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index b64a16d854..24357fdfad 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -1,14 +1,11 @@
 use super::{
     config::TEST_CONFIG,
-    helpers::{
-        self, apply_settings_from_relay_query, connect_and_wait, disconnect_and_wait,
-        set_relay_settings,
-    },
+    helpers::{self, apply_settings_from_relay_query, connect_and_wait, disconnect_and_wait},
     Error, TestContext,
 };
 use crate::{
     network_monitor::{start_packet_monitor, MonitorOptions},
-    tests::helpers::login_with_retries,
+    tests::helpers::{login_with_retries, update_relay_constraints},
 };
 
 use anyhow::Context;
@@ -17,8 +14,7 @@ use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
-        self, BridgeConstraints, BridgeSettings, BridgeType, OpenVpnConstraints, RelayConstraints,
-        RelaySettings, TransportPort,
+        self, BridgeConstraints, BridgeSettings, BridgeType, OpenVpnConstraints, TransportPort,
     },
     wireguard,
 };
@@ -63,15 +59,12 @@ pub async fn test_openvpn_tunnel(
     for (protocol, constraint) in CONSTRAINTS {
         log::info!("Connect to {protocol} OpenVPN endpoint");
 
-        let relay_settings = RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            openvpn_constraints: OpenVpnConstraints { port: constraint },
-            ..Default::default()
-        });
-
-        set_relay_settings(&mut mullvad_client, relay_settings)
-            .await
-            .expect("failed to update relay settings");
+        update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+            relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+            relay_constraints.openvpn_constraints = OpenVpnConstraints { port: constraint };
+        })
+        .await
+        .expect("failed to update relay constraints");
 
         connect_and_wait(&mut mullvad_client).await?;
 
@@ -353,15 +346,11 @@ pub async fn test_wireguard_autoconnect(
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
     log::info!("Setting tunnel protocol to WireGuard");
-
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::Wireguard);
+    })
+    .await
+    .expect("failed to update relay constraints");
 
     mullvad_client
         .set_auto_connect(true)
@@ -396,14 +385,11 @@ pub async fn test_openvpn_autoconnect(
 ) -> Result<(), Error> {
     log::info!("Setting tunnel protocol to OpenVPN");
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
+    .await
+    .expect("failed to update relay constraints");
 
     mullvad_client
         .set_auto_connect(true)
@@ -596,15 +582,11 @@ pub async fn test_remote_socks_bridge(
         .await
         .expect("failed to update bridge settings");
 
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            ..Default::default()
-        }),
-    )
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
     .await
-    .expect("failed to update relay settings");
+    .expect("failed to update relay constraints");
 
     // Connect to VPN
     //
@@ -694,15 +676,11 @@ pub async fn test_local_socks_bridge(
         .await
         .expect("failed to update bridge settings");
 
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            ..Default::default()
-        }),
-    )
+    update_relay_constraints(&mut mullvad_client, |relay_constraints| {
+        relay_constraints.tunnel_protocol = Constraint::Only(TunnelType::OpenVpn);
+    })
     .await
-    .expect("failed to update relay settings");
+    .expect("failed to update relay constraints");
 
     // Connect to VPN
     //
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index 5ccf0f863e..4005f76480 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -1,19 +1,20 @@
 use super::{
     helpers::{
-        self, connect_and_wait, send_guest_probes, set_relay_settings,
-        unreachable_wireguard_tunnel, wait_for_tunnel_state,
+        self, connect_and_wait, send_guest_probes, unreachable_wireguard_tunnel,
+        wait_for_tunnel_state,
     },
     ui, Error, TestContext,
 };
-use crate::{assert_tunnel_state, tests::helpers::ping_sized_with_timeout};
+use crate::{
+    assert_tunnel_state,
+    tests::helpers::{ping_sized_with_timeout, set_custom_endpoint, update_relay_constraints},
+};
 
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use mullvad_types::{
     constraints::Constraint,
-    relay_constraints::{
-        GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
-    },
+    relay_constraints::{GeographicLocationConstraint, LocationConstraint},
     states::TunnelState,
     CustomTunnelEndpoint,
 };
@@ -185,14 +186,15 @@ pub async fn test_connecting_state(
     log::info!("Verify tunnel state: disconnected");
     assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected { .. });
 
-    let relay_settings = RelaySettings::CustomTunnelEndpoint(CustomTunnelEndpoint {
-        host: "1.3.3.7".to_owned(),
-        config: mullvad_types::ConnectionConfig::Wireguard(unreachable_wireguard_tunnel()),
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    set_custom_endpoint(
+        &mut mullvad_client,
+        CustomTunnelEndpoint {
+            host: "1.3.3.7".to_owned(),
+            config: mullvad_types::ConnectionConfig::Wireguard(unreachable_wireguard_tunnel()),
+        },
+    )
+    .await
+    .expect("failed to update relay settings");
 
     mullvad_client
         .connect_tunnel()
@@ -273,21 +275,18 @@ pub async fn test_error_state(
 
     log::info!("Enter error state");
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        location: Constraint::Only(LocationConstraint::from(
-            GeographicLocationConstraint::country("xx"),
-        )),
-        ..Default::default()
-    });
-
     mullvad_client
         .set_allow_lan(false)
         .await
         .expect("failed to disable LAN sharing");
 
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    update_relay_constraints(&mut mullvad_client, |constraints| {
+        constraints.location = Constraint::Only(LocationConstraint::from(
+            GeographicLocationConstraint::country("xx"),
+        ))
+    })
+    .await
+    .expect("Failed to set invalid location");
 
     let _ = connect_and_wait(&mut mullvad_client).await;
     assert_tunnel_state!(&mut mullvad_client, TunnelState::Error { .. });