Replace footgunny From implementations for RelayQuery

6 files changed, 179 insertions, 195 deletions

diff --git a/mullvad-relay-selector/src/relay_selector/mod.rs b/mullvad-relay-selector/src/relay_selector/mod.rs
index d7efa89b5f..dc495b05aa 100644
--- a/mullvad-relay-selector/src/relay_selector/mod.rs
+++ b/mullvad-relay-selector/src/relay_selector/mod.rs
@@ -369,14 +369,14 @@ impl<'a> TryFrom<NormalSelectorConfig<'a>> for RelayQuery {
                 bridge_settings: match bridge_state {
                     BridgeState::On => match bridge_settings.bridge_type {
                         mullvad_types::relay_constraints::BridgeType::Normal => {
-                            Constraint::Only(BridgeQuery::Normal(bridge_settings.normal.clone()))
+                            BridgeQuery::Normal(bridge_settings.normal.clone())
                         }
                         mullvad_types::relay_constraints::BridgeType::Custom => {
-                            Constraint::Only(BridgeQuery::Custom(bridge_settings.custom.clone()))
+                            BridgeQuery::Custom(bridge_settings.custom.clone())
                         }
                     },
-                    BridgeState::Auto => Constraint::Only(BridgeQuery::Auto),
-                    BridgeState::Off => Constraint::Only(BridgeQuery::Off),
+                    BridgeState::Auto => BridgeQuery::Auto,
+                    BridgeState::Off => BridgeQuery::Off,
                 },
             }
         }
@@ -930,7 +930,7 @@ impl RelaySelector {
         if !BridgeQuery::should_use_bridge(&query.openvpn_constraints().bridge_settings) {
             Ok(None)
         } else {
-            let bridge_query = &query.openvpn_constraints().bridge_settings.clone().unwrap();
+            let bridge_query = &query.openvpn_constraints().bridge_settings;
             let custom_lists = &custom_lists;
             match protocol {
                 TransportProtocol::Udp => {
diff --git a/mullvad-relay-selector/src/relay_selector/query.rs b/mullvad-relay-selector/src/relay_selector/query.rs
index 08e75d782a..67fe43e2a3 100644
--- a/mullvad-relay-selector/src/relay_selector/query.rs
+++ b/mullvad-relay-selector/src/relay_selector/query.rs
@@ -32,9 +32,10 @@ use crate::{AdditionalWireguardConstraints, Error};
 use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
-        BridgeConstraints, LocationConstraint, ObfuscationSettings, OpenVpnConstraints, Ownership,
-        Providers, RelayConstraints, RelaySettings, SelectedObfuscation, ShadowsocksSettings,
-        TransportPort, Udp2TcpObfuscationSettings, WireguardConstraints,
+        BridgeConstraints, BridgeSettings, BridgeState, BridgeType, LocationConstraint,
+        ObfuscationSettings, OpenVpnConstraints, Ownership, Providers, RelayConstraints,
+        SelectedObfuscation, ShadowsocksSettings, TransportPort, Udp2TcpObfuscationSettings,
+        WireguardConstraints,
     },
     wireguard::QuantumResistantState,
     Intersection,
@@ -179,6 +180,42 @@ impl RelayQuery {
 
         Ok(())
     }
+
+    /// The mapping from [`RelayQuery`] to all underlying settings types.
+    ///
+    /// Useful in contexts where you cannot use the query directly but
+    /// still want use of the builder for convenience. For example in
+    /// end to end tests where you must use the management interface
+    /// to apply settings to the daemon.
+    pub fn into_settings(
+        self,
+    ) -> (
+        RelayConstraints,
+        BridgeState,
+        BridgeSettings,
+        ObfuscationSettings,
+    ) {
+        let (bridge_state, bridge_settings) = self
+            .openvpn_constraints
+            .bridge_settings
+            .clone()
+            .into_settings();
+        let obfuscation = self
+            .wireguard_constraints
+            .obfuscation
+            .clone()
+            .into_settings();
+        let constraints = RelayConstraints {
+            location: self.location,
+            providers: self.providers,
+            ownership: self.ownership,
+            tunnel_protocol: self.tunnel_protocol,
+            wireguard_constraints: self.wireguard_constraints.into_constraints(),
+            openvpn_constraints: self.openvpn_constraints.into_constraints(),
+        };
+
+        (constraints, bridge_state, bridge_settings, obfuscation)
+    }
 }
 
 impl Default for RelayQuery {
@@ -207,26 +244,6 @@ impl Default for RelayQuery {
     }
 }
 
-impl From<RelayQuery> for RelayConstraints {
-    /// The mapping from [`RelayQuery`] to [`RelayConstraints`].
-    fn from(value: RelayQuery) -> Self {
-        RelayConstraints {
-            location: value.location,
-            providers: value.providers,
-            ownership: value.ownership,
-            tunnel_protocol: value.tunnel_protocol,
-            wireguard_constraints: WireguardConstraints::from(value.wireguard_constraints),
-            openvpn_constraints: OpenVpnConstraints::from(value.openvpn_constraints),
-        }
-    }
-}
-
-impl From<RelayQuery> for RelaySettings {
-    fn from(value: RelayQuery) -> Self {
-        RelayConstraints::from(value).into()
-    }
-}
-
 /// A query for a relay with Wireguard-specific properties, such as `multihop` and [wireguard
 /// obfuscation][`SelectedObfuscation`].
 ///
@@ -256,6 +273,31 @@ pub enum ObfuscationQuery {
     Shadowsocks(ShadowsocksSettings),
 }
 
+impl ObfuscationQuery {
+    fn into_settings(self) -> ObfuscationSettings {
+        match self {
+            ObfuscationQuery::Auto => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Auto,
+                ..Default::default()
+            },
+            ObfuscationQuery::Off => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Off,
+                ..Default::default()
+            },
+            ObfuscationQuery::Udp2tcp(settings) => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Udp2Tcp,
+                udp2tcp: settings,
+                ..Default::default()
+            },
+            ObfuscationQuery::Shadowsocks(settings) => ObfuscationSettings {
+                selected_obfuscation: SelectedObfuscation::Shadowsocks,
+                shadowsocks: settings,
+                ..Default::default()
+            },
+        }
+    }
+}
+
 impl From<ObfuscationSettings> for ObfuscationQuery {
     /// A query for obfuscation settings.
     ///
@@ -307,6 +349,16 @@ impl WireguardRelayQuery {
             quantum_resistant: QuantumResistantState::Auto,
         }
     }
+
+    /// The mapping from [`WireguardRelayQuery`] to [`WireguardConstraints`].
+    fn into_constraints(self) -> WireguardConstraints {
+        WireguardConstraints {
+            port: self.port,
+            ip_version: self.ip_version,
+            entry_location: self.entry_location,
+            use_multihop: self.use_multihop.unwrap_or(false),
+        }
+    }
 }
 
 impl Default for WireguardRelayQuery {
@@ -315,18 +367,6 @@ impl Default for WireguardRelayQuery {
     }
 }
 
-impl From<WireguardRelayQuery> for WireguardConstraints {
-    /// The mapping from [`WireguardRelayQuery`] to [`WireguardConstraints`].
-    fn from(value: WireguardRelayQuery) -> Self {
-        WireguardConstraints {
-            port: value.port,
-            ip_version: value.ip_version,
-            entry_location: value.entry_location,
-            use_multihop: value.use_multihop.unwrap_or(false),
-        }
-    }
-}
-
 impl From<WireguardRelayQuery> for AdditionalWireguardConstraints {
     /// The mapping from [`WireguardRelayQuery`] to [`AdditionalWireguardConstraints`].
     fn from(value: WireguardRelayQuery) -> Self {
@@ -348,16 +388,21 @@ impl From<WireguardRelayQuery> for AdditionalWireguardConstraints {
 #[derive(Debug, Clone, Eq, PartialEq, Intersection)]
 pub struct OpenVpnRelayQuery {
     pub port: Constraint<TransportPort>,
-    pub bridge_settings: Constraint<BridgeQuery>,
+    pub bridge_settings: BridgeQuery,
 }
 
 impl OpenVpnRelayQuery {
     pub const fn new() -> OpenVpnRelayQuery {
         OpenVpnRelayQuery {
             port: Constraint::Any,
-            bridge_settings: Constraint::Any,
+            bridge_settings: BridgeQuery::Auto,
         }
     }
+
+    /// The mapping from [`OpenVpnRelayQuery`] to [`OpenVpnConstraints`].
+    fn into_constraints(self) -> OpenVpnConstraints {
+        OpenVpnConstraints { port: self.port }
+    }
 }
 
 impl Default for OpenVpnRelayQuery {
@@ -371,7 +416,7 @@ impl Default for OpenVpnRelayQuery {
 ///
 /// [`BridgeState`]: mullvad_types::relay_constraints::BridgeState
 /// [`BridgeSettings`]: mullvad_types::relay_constraints::BridgeSettings
-#[derive(Debug, Clone, Eq, PartialEq)]
+#[derive(Debug, Default, Clone, Eq, PartialEq)]
 pub enum BridgeQuery {
     /// Bridges should not be used.
     Off,
@@ -379,6 +424,7 @@ pub enum BridgeQuery {
     ///
     /// If this variant is intersected with another [`BridgeQuery`] `bq`,
     /// `bq` is always preferred.
+    #[default]
     Auto,
     /// Bridges should be used.
     Normal(BridgeConstraints),
@@ -387,22 +433,38 @@ pub enum BridgeQuery {
 }
 
 impl BridgeQuery {
-    /// If `bridge_constraints` is `Any`, bridges should not be used due to
-    /// latency concerns.
-    ///
-    /// If `bridge_constraints` is `Only(settings)`, then `settings` will be
-    /// used to decide if bridges should be used. See [`BridgeQuery`] for more
-    /// details, but the algorithm beaks down to this:
+    /// `settings` will be used to decide if bridges should be used. See [`BridgeQuery`]
+    /// for more details, but the algorithm beaks down to this:
     ///
     /// * `BridgeQuery::Off`: bridges will not be used
     /// * otherwise: bridges should be used
-    pub const fn should_use_bridge(bridge_constraints: &Constraint<BridgeQuery>) -> bool {
-        match bridge_constraints {
-            Constraint::Only(settings) => match settings {
-                BridgeQuery::Normal(_) | BridgeQuery::Custom(_) => true,
-                BridgeQuery::Off | BridgeQuery::Auto => false,
-            },
-            Constraint::Any => false,
+    pub const fn should_use_bridge(settings: &BridgeQuery) -> bool {
+        match settings {
+            BridgeQuery::Normal(_) | BridgeQuery::Custom(_) => true,
+            BridgeQuery::Off | BridgeQuery::Auto => false,
+        }
+    }
+
+    fn into_settings(self) -> (BridgeState, BridgeSettings) {
+        match self {
+            BridgeQuery::Off => (BridgeState::Off, Default::default()),
+            BridgeQuery::Auto => (BridgeState::Auto, Default::default()),
+            BridgeQuery::Normal(constraints) => (
+                BridgeState::On,
+                BridgeSettings {
+                    bridge_type: BridgeType::Normal,
+                    normal: constraints,
+                    custom: None,
+                },
+            ),
+            BridgeQuery::Custom(custom) => (
+                BridgeState::On,
+                BridgeSettings {
+                    bridge_type: BridgeType::Normal,
+                    normal: Default::default(),
+                    custom,
+                },
+            ),
         }
     }
 }
@@ -425,13 +487,6 @@ impl Intersection for BridgeQuery {
     }
 }
 
-impl From<OpenVpnRelayQuery> for OpenVpnConstraints {
-    /// The mapping from [`OpenVpnRelayQuery`] to [`OpenVpnConstraints`].
-    fn from(value: OpenVpnRelayQuery) -> Self {
-        OpenVpnConstraints { port: value.port }
-    }
-}
-
 #[allow(unused)]
 pub mod builder {
     //! Strongly typed Builder pattern for of relay constraints though the use of the Typestate
@@ -498,10 +553,6 @@ pub mod builder {
             debug_assert!(self.query.validate().is_ok());
             self.query
         }
-
-        pub fn into_constraint(self) -> RelayConstraints {
-            RelayConstraints::from(self.build())
-        }
     }
 
     impl RelayQueryBuilder<Any> {
@@ -802,8 +853,7 @@ pub mod builder {
                 bridge_settings: bridge_settings.clone(),
             };
 
-            self.query.openvpn_constraints.bridge_settings =
-                Constraint::Only(BridgeQuery::Normal(bridge_settings));
+            self.query.openvpn_constraints.bridge_settings = BridgeQuery::Normal(bridge_settings);
 
             let builder = RelayQueryBuilder {
                 query: self.query,
@@ -820,14 +870,14 @@ pub mod builder {
             self.protocol.bridge_settings.location =
                 Constraint::Only(LocationConstraint::from(location));
             self.query.openvpn_constraints.bridge_settings =
-                Constraint::Only(BridgeQuery::Normal(self.protocol.bridge_settings.clone()));
+                BridgeQuery::Normal(self.protocol.bridge_settings.clone());
             self
         }
         /// Constrain the [`Providers`] of the selected bridge.
         pub fn bridge_providers(mut self, providers: Providers) -> Self {
             self.protocol.bridge_settings.providers = Constraint::Only(providers);
             self.query.openvpn_constraints.bridge_settings =
-                Constraint::Only(BridgeQuery::Normal(self.protocol.bridge_settings.clone()));
+                BridgeQuery::Normal(self.protocol.bridge_settings.clone());
             self
         }
         /// Constrain the [`Ownership`] of the selected bridge.
diff --git a/mullvad-relay-selector/tests/relay_selector.rs b/mullvad-relay-selector/tests/relay_selector.rs
index 51a395f6d5..bdfd8e19d0 100644
--- a/mullvad-relay-selector/tests/relay_selector.rs
+++ b/mullvad-relay-selector/tests/relay_selector.rs
@@ -1292,7 +1292,7 @@ fn openvpn_handle_bridge_settings() {
     // should fail.
     query
         .set_openvpn_constraints(OpenVpnRelayQuery {
-            bridge_settings: Constraint::Only(BridgeQuery::Normal(BridgeConstraints::default())),
+            bridge_settings: BridgeQuery::Normal(BridgeConstraints::default()),
             ..query.openvpn_constraints().clone()
         })
         .unwrap();
@@ -1524,9 +1524,10 @@ fn valid_user_setting_should_yield_relay() {
     // Make a valid user relay constraint
     let location = GeographicLocationConstraint::hostname("se", "got", "se9-wireguard");
     let user_query = RelayQueryBuilder::new().location(location.clone()).build();
-    let user_constraints = RelayQueryBuilder::new()
+    let (user_constraints, ..) = RelayQueryBuilder::new()
         .location(location.clone())
-        .into_constraint();
+        .build()
+        .into_settings();
 
     let config = SelectorConfig {
         relay_settings: user_constraints.into(),
diff --git a/mullvad-types/src/constraints/constraint.rs b/mullvad-types/src/constraints/constraint.rs
index 35b83e1320..2fb2f163c2 100644
--- a/mullvad-types/src/constraints/constraint.rs
+++ b/mullvad-types/src/constraints/constraint.rs
@@ -41,6 +41,15 @@ impl<T: fmt::Display> fmt::Display for Constraint<T> {
     }
 }
 
+impl<T: Default> Constraint<T> {
+    pub fn unwrap_or_default(self) -> T {
+        match self {
+            Constraint::Any => Default::default(),
+            Constraint::Only(value) => value,
+        }
+    }
+}
+
 impl<T> Constraint<T> {
     pub fn unwrap(self) -> T {
         match self {
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 542cf0c8a4..6c51b1f185 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -19,8 +19,7 @@ use mullvad_types::{
     constraints::Constraint,
     location::Location,
     relay_constraints::{
-        BridgeSettings, GeographicLocationConstraint, LocationConstraint, OpenVpnConstraints,
-        RelayConstraints, RelaySettings, WireguardConstraints,
+        GeographicLocationConstraint, LocationConstraint, RelayConstraints, RelaySettings,
     },
     relay_list::Relay,
     states::TunnelState,
@@ -589,24 +588,38 @@ impl<T> Drop for AbortOnDrop<T> {
     }
 }
 
-pub async fn set_relay_settings(
+pub async fn apply_settings_from_relay_query(
     mullvad_client: &mut MullvadProxyClient,
-    relay_settings: impl Into<RelaySettings>,
+    query: RelayQuery,
 ) -> Result<(), Error> {
+    let (constraints, bridge_state, bridge_settings, obfuscation) = query.into_settings();
+
     mullvad_client
-        .set_relay_settings(relay_settings.into())
+        .set_relay_settings(constraints.into())
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))
+        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))?;
+    mullvad_client
+        .set_bridge_state(bridge_state)
+        .await
+        .map_err(|error| Error::Daemon(format!("Failed to set bridge state: {}", error)))?;
+    mullvad_client
+        .set_bridge_settings(bridge_settings)
+        .await
+        .map_err(|error| Error::Daemon(format!("Failed to set bridge settings: {}", error)))?;
+    mullvad_client
+        .set_obfuscation_settings(obfuscation)
+        .await
+        .map_err(|error| Error::Daemon(format!("Failed to set obfuscation settings: {}", error)))
 }
 
-pub async fn set_bridge_settings(
+pub async fn set_relay_settings(
     mullvad_client: &mut MullvadProxyClient,
-    bridge_settings: BridgeSettings,
+    relay_settings: impl Into<RelaySettings>,
 ) -> Result<(), Error> {
     mullvad_client
-        .set_bridge_settings(bridge_settings)
+        .set_relay_settings(relay_settings.into())
         .await
-        .map_err(|error| Error::Daemon(format!("Failed to set bridge settings: {}", error)))
+        .map_err(|error| Error::Daemon(format!("Failed to set relay settings: {}", error)))
 }
 
 /// Wait for the relay list to be updated, to make sure we have the overridden one.
@@ -698,16 +711,8 @@ pub async fn constrain_to_relay(
             }
             | GetRelay::OpenVpn { exit, .. } => {
                 let location = into_constraint(&exit)?;
-                let relay_constraints = RelayConstraints {
-                    location,
-                    wireguard_constraints: WireguardConstraints::from(
-                        query.wireguard_constraints().clone(),
-                    ),
-                    openvpn_constraints: OpenVpnConstraints::from(
-                        query.openvpn_constraints().clone(),
-                    ),
-                    ..Default::default()
-                };
+                let (mut relay_constraints, ..) = query.into_settings();
+                relay_constraints.location = location;
                 Ok((exit, relay_constraints))
             }
             unsupported => bail!("Can not constrain to a {unsupported:?}"),
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index f8ba47d1ca..575339c5a8 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -1,7 +1,8 @@
 use super::{
     config::TEST_CONFIG,
     helpers::{
-        self, connect_and_wait, disconnect_and_wait, set_bridge_settings, set_relay_settings,
+        self, apply_settings_from_relay_query, connect_and_wait, disconnect_and_wait,
+        set_relay_settings,
     },
     Error, TestContext,
 };
@@ -17,8 +18,7 @@ use mullvad_types::{
     constraints::Constraint,
     relay_constraints::{
         self, BridgeConstraints, BridgeSettings, BridgeType, OpenVpnConstraints, RelayConstraints,
-        RelaySettings, SelectedObfuscation, ShadowsocksSettings, TransportPort,
-        Udp2TcpObfuscationSettings, WireguardConstraints,
+        RelaySettings, TransportPort, WireguardConstraints,
     },
     states::TunnelState,
     wireguard,
@@ -145,25 +145,9 @@ pub async fn test_udp2tcp_tunnel(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
-    mullvad_client
-        .set_obfuscation_settings(relay_constraints::ObfuscationSettings {
-            selected_obfuscation: SelectedObfuscation::Udp2Tcp,
-            udp2tcp: Udp2TcpObfuscationSettings {
-                port: Constraint::Any,
-            },
-            ..Default::default()
-        })
-        .await
-        .expect("failed to enable udp2tcp");
-
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-        ..Default::default()
-    });
+    let query = RelayQueryBuilder::new().wireguard().udp2tcp().build();
 
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .expect("failed to update relay settings");
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     log::info!("Connect to WireGuard via tcp2udp endpoint");
 
@@ -211,25 +195,9 @@ pub async fn test_wireguard_over_shadowsocks(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    mullvad_client
-        .set_obfuscation_settings(relay_constraints::ObfuscationSettings {
-            selected_obfuscation: SelectedObfuscation::Shadowsocks,
-            shadowsocks: ShadowsocksSettings {
-                port: Constraint::Any,
-            },
-            ..Default::default()
-        })
-        .await
-        .context("Failed to enable shadowsocks")?;
+    let query = RelayQueryBuilder::new().wireguard().shadowsocks().build();
 
-    let relay_settings = RelaySettings::Normal(RelayConstraints {
-        tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
-        ..Default::default()
-    });
-
-    set_relay_settings(&mut mullvad_client, relay_settings)
-        .await
-        .context("Failed to update relay settings")?;
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     log::info!("Connect to WireGuard via shadowsocks endpoint");
 
@@ -259,24 +227,9 @@ pub async fn test_bridge(
     //
     log::info!("Updating bridge settings");
 
-    mullvad_client
-        .set_bridge_state(relay_constraints::BridgeState::On)
-        .await
-        .expect("failed to enable bridge mode");
+    let query = RelayQueryBuilder::new().openvpn().bridge().build();
 
-    set_bridge_settings(&mut mullvad_client, BridgeSettings::default())
-        .await
-        .expect("failed to update bridge settings");
-
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(RelayConstraints {
-            tunnel_protocol: Constraint::Only(TunnelType::OpenVpn),
-            ..Default::default()
-        }),
-    )
-    .await
-    .expect("failed to update relay settings");
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     // Connect to VPN
     //
@@ -341,17 +294,9 @@ pub async fn test_multihop(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> Result<(), Error> {
-    let relay_constraints = RelayQueryBuilder::new()
-        .wireguard()
-        .multihop()
-        .into_constraint();
+    let query = RelayQueryBuilder::new().wireguard().multihop().build();
 
-    set_relay_settings(
-        &mut mullvad_client,
-        RelaySettings::Normal(relay_constraints),
-    )
-    .await
-    .expect("failed to update relay settings");
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     // Connect
     //
@@ -458,7 +403,7 @@ pub async fn test_daita(
 ) -> anyhow::Result<()> {
     log::info!("Connecting to relay with DAITA");
 
-    set_relay_settings(
+    apply_settings_from_relay_query(
         &mut mullvad_client,
         RelayQueryBuilder::new().wireguard().build(),
     )
@@ -616,26 +561,13 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
         .await
         .expect("Failed to enable PQ tunnels");
 
-    mullvad_client
-        .set_obfuscation_settings(relay_constraints::ObfuscationSettings {
-            selected_obfuscation: SelectedObfuscation::Udp2Tcp,
-            udp2tcp: Udp2TcpObfuscationSettings {
-                port: Constraint::Any,
-            },
-            ..Default::default()
-        })
-        .await
-        .expect("Failed to enable obfuscation");
-
-    let relay_constraints = RelayQueryBuilder::new()
+    let query = RelayQueryBuilder::new()
         .wireguard()
         .multihop()
-        .into_constraint();
+        .udp2tcp()
+        .build();
 
-    mullvad_client
-        .set_relay_settings(RelaySettings::Normal(relay_constraints))
-        .await
-        .expect("Failed to update relay settings");
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     connect_and_wait(&mut mullvad_client).await?;
 
@@ -664,26 +596,13 @@ pub async fn test_quantum_resistant_multihop_shadowsocks_tunnel(
         .await
         .context("Failed to enable PQ tunnels")?;
 
-    mullvad_client
-        .set_obfuscation_settings(relay_constraints::ObfuscationSettings {
-            selected_obfuscation: SelectedObfuscation::Shadowsocks,
-            shadowsocks: ShadowsocksSettings {
-                port: Constraint::Any,
-            },
-            ..Default::default()
-        })
-        .await
-        .context("Failed to enable obfuscation")?;
-
-    let relay_constraints = RelayQueryBuilder::new()
+    let query = RelayQueryBuilder::new()
         .wireguard()
         .multihop()
-        .into_constraint();
+        .shadowsocks()
+        .build();
 
-    mullvad_client
-        .set_relay_settings(RelaySettings::Normal(relay_constraints))
-        .await
-        .context("Failed to update relay settings")?;
+    apply_settings_from_relay_query(&mut mullvad_client, query).await?;
 
     connect_and_wait(&mut mullvad_client).await?;