test-framework: Prefer types from `mullvad-types`

Prefer types from `mullvad-types` over equivavelent types defined in
`mullvad-management-interface`.

5 files changed, 101 insertions, 64 deletions

diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index e7c2e12d0f..74b07de44e 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -3,11 +3,13 @@ use crate::network_monitor::{start_packet_monitor, MonitorOptions};
 use futures::StreamExt;
 use mullvad_management_interface::{types, ManagementServiceClient};
 use mullvad_types::{
+    location::Location,
     relay_constraints::{
-        BridgeState, Constraint, GeographicLocationConstraint, LocationConstraint,
+        BridgeSettings, BridgeState, Constraint, GeographicLocationConstraint, LocationConstraint,
         ObfuscationSettings, OpenVpnConstraints, RelayConstraints, RelaySettings,
         WireguardConstraints,
     },
+    relay_list::{Relay, RelayList},
     states::TunnelState,
 };
 use pnet_packet::ip::IpNextHeaderProtocols;
@@ -378,6 +380,19 @@ pub async fn set_relay_settings(
     Ok(())
 }
 
+pub async fn set_bridge_settings(
+    mullvad_client: &mut ManagementServiceClient,
+    bridge_settings: BridgeSettings,
+) -> Result<(), Error> {
+    let new_settings = types::BridgeSettings::from(bridge_settings);
+
+    mullvad_client
+        .set_bridge_settings(new_settings)
+        .await
+        .map_err(|error| Error::DaemonError(format!("Failed to set bridge settings: {}", error)))?;
+    Ok(())
+}
+
 pub async fn get_tunnel_state(mullvad_client: &mut ManagementServiceClient) -> TunnelState {
     let state = mullvad_client
         .get_tunnel_state(())
@@ -434,6 +449,25 @@ pub fn unreachable_wireguard_tunnel() -> talpid_types::net::wireguard::Connectio
     }
 }
 
+/// Find a relay from the daemon's relay list that matches `critera`.
+///
+/// * `mullvad_client` - An interface to the Mullvad daemon.
+/// * `critera` - A function used to determine which relays to include in random selection.
+pub async fn relay<Filter>(
+    mullvad_client: &mut ManagementServiceClient,
+    criteria: Filter,
+) -> Result<Relay, Error>
+where
+    Filter: Fn(&Relay) -> bool,
+{
+    filter_relays(mullvad_client, criteria)
+        .await?
+        .pop()
+        .ok_or(Error::Other(
+            "No mathing bridge was found in the relay list".to_string(),
+        ))
+}
+
 /// Randomly select an entry and exit node from the daemon's relay list.
 /// The exit node is distinct from the entry node.
 ///
@@ -442,9 +476,9 @@ pub fn unreachable_wireguard_tunnel() -> talpid_types::net::wireguard::Connectio
 pub async fn random_entry_and_exit<Filter>(
     mullvad_client: &mut ManagementServiceClient,
     criteria: Filter,
-) -> Result<(types::Relay, types::Relay), Error>
+) -> Result<(Relay, Relay), Error>
 where
-    Filter: Fn(&types::Relay) -> bool,
+    Filter: Fn(&Relay) -> bool,
 {
     use itertools::Itertools;
     // Pluck the first 2 relays and return them as a tuple.
@@ -465,43 +499,35 @@ where
 pub async fn filter_relays<Filter>(
     mullvad_client: &mut ManagementServiceClient,
     criteria: Filter,
-) -> Result<Vec<types::Relay>, Error>
+) -> Result<Vec<Relay>, Error>
 where
-    Filter: Fn(&types::Relay) -> bool,
+    Filter: Fn(&Relay) -> bool,
 {
-    let relaylist = mullvad_client
+    let relay_list: RelayList = mullvad_client
         .get_relay_locations(())
         .await
         .map_err(|error| Error::DaemonError(format!("Failed to obtain relay list: {}", error)))?
-        .into_inner();
+        .into_inner()
+        .try_into()?;
 
-    Ok(flatten_relaylist(relaylist)
-        .into_iter()
-        .filter(criteria)
+    Ok(relay_list
+        .relays()
+        .filter(|relay| criteria(relay))
+        .cloned()
         .collect())
 }
 
-/// Dig out the [`Relay`]s contained in a [`RelayList`].
-pub fn flatten_relaylist(relays: types::RelayList) -> Vec<types::Relay> {
-    relays
-        .countries
-        .iter()
-        .flat_map(|country| country.cities.clone())
-        .flat_map(|city| city.relays)
-        .collect()
-}
-
 /// Convenience function for constructing a constraint from a given [`Relay`].
 ///
 /// # Panics
 ///
 /// The relay must have a location set.
-pub fn into_constraint(relay: &types::Relay) -> Constraint<LocationConstraint> {
+pub fn into_constraint(relay: &Relay) -> Constraint<LocationConstraint> {
     relay
         .location
         .as_ref()
         .map(
-            |types::Location {
+            |Location {
                  country_code,
                  city_code,
                  ..
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index e96e4ad0ba..b35f1b9b7a 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -17,7 +17,10 @@ use test_rpc::ServiceClient;
 
 use futures::future::BoxFuture;
 
-use mullvad_management_interface::{types::Settings, ManagementServiceClient};
+use mullvad_management_interface::{
+    types::{self, Settings},
+    ManagementServiceClient,
+};
 use once_cell::sync::OnceCell;
 use std::time::Duration;
 
@@ -37,7 +40,7 @@ pub type TestWrapperFunction = Box<
     ) -> BoxFuture<'static, Result<(), Error>>,
 >;
 
-#[derive(err_derive::Error, Debug, PartialEq, Eq)]
+#[derive(err_derive::Error, Debug)]
 pub enum Error {
     #[error(display = "RPC call failed")]
     Rpc(#[source] test_rpc::Error),
@@ -57,6 +60,9 @@ pub enum Error {
     #[error(display = "The daemon returned an error: {}", _0)]
     DaemonError(String),
 
+    #[error(display = "Failed to parse gRPC response")]
+    InvalidGrpcResponse(#[error(source)] types::FromProtobufTypeError),
+
     #[error(display = "An error occurred: {}", _0)]
     Other(String),
 }
diff --git a/test/test-manager/src/tests/tunnel.rs b/test/test-manager/src/tests/tunnel.rs
index 1592f7c2a3..bedba2752f 100644
--- a/test/test-manager/src/tests/tunnel.rs
+++ b/test/test-manager/src/tests/tunnel.rs
@@ -1,6 +1,7 @@
-use super::helpers::{self, connect_and_wait, disconnect_and_wait, set_relay_settings};
+use super::helpers::{
+    self, connect_and_wait, disconnect_and_wait, set_bridge_settings, set_relay_settings,
+};
 use super::{Error, TestContext};
-use std::net::IpAddr;
 
 use crate::network_monitor::{start_packet_monitor, MonitorOptions};
 use mullvad_management_interface::{types, ManagementServiceClient};
@@ -9,6 +10,7 @@ use mullvad_types::relay_constraints::{
     OpenVpnConstraints, RelayConstraints, RelaySettings, SelectedObfuscation, TransportPort,
     Udp2TcpObfuscationSettings, WireguardConstraints,
 };
+use mullvad_types::relay_list::{Relay, RelayEndpointData};
 use mullvad_types::wireguard;
 use pnet_packet::ip::IpNextHeaderProtocols;
 use talpid_types::net::{TransportProtocol, TunnelType};
@@ -200,21 +202,22 @@ pub async fn test_bridge(
     rpc: ServiceClient,
     mut mullvad_client: ManagementServiceClient,
 ) -> Result<(), Error> {
-    log::info!("Select relay");
-    let bridge_filter = |bridge: &types::Relay| {
-        bridge.active && bridge.endpoint_type == i32::from(types::relay::RelayType::Bridge)
-    };
-    let ovpn_filter = |relay: &types::Relay| {
-        relay.active && relay.endpoint_type == i32::from(types::relay::RelayType::Openvpn)
-    };
-    let entry = helpers::filter_relays(&mut mullvad_client, bridge_filter)
-        .await?
-        .pop()
-        .unwrap();
-    let exit = helpers::filter_relays(&mut mullvad_client, ovpn_filter)
-        .await?
-        .pop()
-        .unwrap();
+    let entry = helpers::relay(&mut mullvad_client, |bridge| {
+        bridge.active && matches!(bridge.endpoint_data, RelayEndpointData::Bridge)
+    })
+    .await?;
+    let exit = helpers::relay(&mut mullvad_client, |relay| {
+        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Openvpn)
+    })
+    .await?;
+
+    log::info!(
+        "Selected entry bridge {entry}:{entry_ip} & exit relay {exit}:{exit_ip}",
+        entry = entry.hostname,
+        entry_ip = entry.ipv4_addr_in.to_string(),
+        exit = exit.hostname,
+        exit_ip = exit.ipv4_addr_in.to_string()
+    );
 
     //
     // Enable bridge mode
@@ -227,13 +230,12 @@ pub async fn test_bridge(
         .await
         .expect("failed to enable bridge mode");
 
-    mullvad_client
-        .set_bridge_settings(types::BridgeSettings::from(BridgeSettings::Normal(
-            BridgeConstraints {
-                location: helpers::into_constraint(&entry),
-                ..Default::default()
-            },
-        )))
+    let bridge_settings = BridgeSettings::Normal(BridgeConstraints {
+        location: helpers::into_constraint(&entry),
+        ..Default::default()
+    });
+
+    set_bridge_settings(&mut mullvad_client, bridge_settings)
         .await
         .expect("failed to update bridge settings");
 
@@ -254,7 +256,7 @@ pub async fn test_bridge(
     log::info!("Connect to OpenVPN relay via bridge");
 
     let monitor = start_packet_monitor(
-        move |packet| packet.destination.ip() == entry.ipv4_addr_in.parse::<IpAddr>().unwrap(),
+        move |packet| packet.destination.ip() == entry.ipv4_addr_in,
         MonitorOptions::default(),
     )
     .await;
@@ -279,6 +281,8 @@ pub async fn test_bridge(
     // Verify exit IP
     //
 
+    log::info!("Verifying exit server");
+
     assert!(
         helpers::using_mullvad_exit(&rpc).await,
         "expected Mullvad exit IP"
@@ -304,8 +308,8 @@ pub async fn test_multihop(
     //
 
     log::info!("Select relay");
-    let relay_filter = |relay: &types::Relay| {
-        relay.active && relay.endpoint_type == i32::from(types::relay::RelayType::Wireguard)
+    let relay_filter = |relay: &Relay| {
+        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
     };
     let (entry, exit) = helpers::random_entry_and_exit(&mut mullvad_client, relay_filter).await?;
     let exit_constraint = helpers::into_constraint(&exit);
@@ -331,7 +335,7 @@ pub async fn test_multihop(
 
     let monitor = start_packet_monitor(
         move |packet| {
-            packet.destination.ip() == entry.ipv4_addr_in.parse::<IpAddr>().unwrap()
+            packet.destination.ip() == entry.ipv4_addr_in
                 && packet.protocol == IpNextHeaderProtocols::Udp
         },
         MonitorOptions::default(),
@@ -565,9 +569,8 @@ pub async fn test_quantum_resistant_multihop_udp2tcp_tunnel(
                 wireguard_constraints: WireguardConstraints {
                     use_multihop: true,
                     ..Default::default()
-                }
-                .into(),
-                tunnel_protocol: Constraint::Only(TunnelType::Wireguard).into(),
+                },
+                tunnel_protocol: Constraint::Only(TunnelType::Wireguard),
                 ..Default::default()
             },
         )))
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index cf0b7f4efe..87fda3b685 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -6,14 +6,15 @@ use super::{ui, Error, TestContext};
 use crate::assert_tunnel_state;
 use crate::vm::network::DUMMY_LAN_INTERFACE_IP;
 
-use mullvad_management_interface::{types, ManagementServiceClient};
+use mullvad_management_interface::ManagementServiceClient;
 use mullvad_types::relay_constraints::GeographicLocationConstraint;
+use mullvad_types::relay_list::{Relay, RelayEndpointData};
 use mullvad_types::CustomTunnelEndpoint;
 use mullvad_types::{
     relay_constraints::{Constraint, LocationConstraint, RelayConstraints, RelaySettings},
     states::TunnelState,
 };
-use std::net::{IpAddr, Ipv4Addr, SocketAddr};
+use std::net::{IpAddr, SocketAddr};
 use talpid_types::net::{Endpoint, TransportProtocol, TunnelEndpoint, TunnelType};
 use test_macro::test_function;
 use test_rpc::{Interface, ServiceClient};
@@ -257,8 +258,8 @@ pub async fn test_connected_state(
 
     log::info!("Select relay");
 
-    let relay_filter = |relay: &types::Relay| {
-        relay.active && relay.endpoint_type == i32::from(types::relay::RelayType::Wireguard)
+    let relay_filter = |relay: &Relay| {
+        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
     };
 
     let relay = helpers::filter_relays(&mut mullvad_client, relay_filter)
@@ -306,7 +307,7 @@ pub async fn test_connected_state(
                 },
             ..
         } => {
-            assert_eq!(*addr.ip(), relay.ipv4_addr_in.parse::<Ipv4Addr>().unwrap());
+            assert_eq!(*addr.ip(), relay.ipv4_addr_in);
         }
         actual => panic!("unexpected tunnel state: {:?}", actual),
     }
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index 91f124338f..3600eb1d27 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -1,8 +1,9 @@
 use super::config::TEST_CONFIG;
 use super::helpers;
 use super::{Error, TestContext};
-use mullvad_management_interface::{types, ManagementServiceClient};
+use mullvad_management_interface::ManagementServiceClient;
 use mullvad_types::relay_constraints::{RelayConstraints, RelaySettings};
+use mullvad_types::relay_list::{Relay, RelayEndpointData};
 use std::{
     collections::BTreeMap,
     fmt::Debug,
@@ -87,8 +88,8 @@ pub async fn test_ui_tunnel_settings(
 ) -> Result<(), Error> {
     // tunnel-state.spec precondition: a single WireGuard relay should be selected
     log::info!("Select WireGuard relay");
-    let entry = helpers::filter_relays(&mut mullvad_client, |relay: &types::Relay| {
-        relay.active && relay.endpoint_type == i32::from(types::relay::RelayType::Wireguard)
+    let entry = helpers::filter_relays(&mut mullvad_client, |relay: &Relay| {
+        relay.active && matches!(relay.endpoint_data, RelayEndpointData::Wireguard(_))
     })
     .await?
     .pop()
@@ -109,7 +110,7 @@ pub async fn test_ui_tunnel_settings(
         &["tunnel-state.spec"],
         [
             ("HOSTNAME", entry.hostname.as_str()),
-            ("IN_IP", entry.ipv4_addr_in.as_str()),
+            ("IN_IP", &entry.ipv4_addr_in.to_string()),
             (
                 "CONNECTION_CHECK_URL",
                 &format!("https://am.i.{}", TEST_CONFIG.mullvad_host),