Retry DNS lookup for API in test manager

5 files changed, 95 insertions, 58 deletions

diff --git a/test/test-manager/src/tests/account.rs b/test/test-manager/src/tests/account.rs
index 69f605a290..2cc7c5a95e 100644
--- a/test/test-manager/src/tests/account.rs
+++ b/test/test-manager/src/tests/account.rs
@@ -1,13 +1,14 @@
 use crate::tests::helpers::{login_with_retries, THROTTLE_RETRY_DELAY};
 
 use super::{config::TEST_CONFIG, helpers, ui, Error, TestContext};
+use anyhow::Context;
 use mullvad_api::DevicesProxy;
 use mullvad_management_interface::{client::DaemonEvent, MullvadProxyClient};
 use mullvad_types::{
     device::{Device, DeviceState},
     states::TunnelState,
 };
-use std::{net::ToSocketAddrs, time::Duration};
+use std::time::Duration;
 use talpid_types::net::wireguard;
 use test_macro::test_function;
 use test_rpc::ServiceClient;
@@ -18,18 +19,18 @@ pub async fn test_login(
     _: TestContext,
     _rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     // Instruct daemon to log in
     //
 
-    clear_devices(&new_device_client())
+    clear_devices(&new_device_client().await?)
         .await
-        .expect("failed to clear devices");
+        .context("failed to clear devices")?;
 
     log::info!("Logging in/generating device");
     login_with_retries(&mut mullvad_client)
         .await
-        .expect("login failed");
+        .context("login failed")?;
 
     // Wait for the relay list to be updated
     helpers::ensure_updated_relay_list(&mut mullvad_client).await?;
@@ -61,10 +62,10 @@ pub async fn test_too_many_devices(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     log::info!("Using up all devices");
 
-    let device_client = new_device_client();
+    let device_client = new_device_client().await.context("Create device client")?;
 
     const MAX_ATTEMPTS: usize = 15;
 
@@ -131,16 +132,16 @@ pub async fn test_revoked_device(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     log::info!("Logging in/generating device");
     login_with_retries(&mut mullvad_client)
         .await
-        .expect("login failed");
+        .context("login failed")?;
 
     let device_id = mullvad_client
         .get_device()
         .await
-        .expect("failed to get device data")
+        .context("failed to get device data")?
         .into_device()
         .unwrap()
         .device
@@ -150,7 +151,9 @@ pub async fn test_revoked_device(
 
     log::debug!("Removing current device");
 
-    let device_client = new_device_client();
+    let device_client = new_device_client()
+        .await
+        .context("Failed to create device client")?;
     retry_if_throttled(|| {
         device_client.remove(TEST_CONFIG.account_number.clone(), device_id.clone())
     })
@@ -167,7 +170,7 @@ pub async fn test_revoked_device(
     let events = mullvad_client
         .events_listen()
         .await
-        .expect("failed to begin listening for state changes");
+        .context("Failed to begin listening for state changes")?;
     let next_state =
         helpers::find_next_tunnel_state(events, |state| matches!(state, TunnelState::Error(..),));
 
@@ -188,7 +191,7 @@ pub async fn test_revoked_device(
     let device_state = mullvad_client
         .get_device()
         .await
-        .expect("failed to get device data");
+        .context("Failed to get device data")?;
     assert!(
         matches!(device_state, DeviceState::Revoked),
         "expected device to be revoked"
@@ -202,7 +205,7 @@ pub async fn test_revoked_device(
 }
 
 /// Remove all devices on the current account
-pub async fn clear_devices(device_client: &DevicesProxy) -> Result<(), mullvad_api::rest::Error> {
+pub async fn clear_devices(device_client: &DevicesProxy) -> anyhow::Result<()> {
     log::info!("Removing all devices for account");
 
     for dev in list_devices_with_retries(device_client).await?.into_iter() {
@@ -216,16 +219,16 @@ pub async fn clear_devices(device_client: &DevicesProxy) -> Result<(), mullvad_a
     Ok(())
 }
 
-pub fn new_device_client() -> DevicesProxy {
+pub async fn new_device_client() -> anyhow::Result<DevicesProxy> {
     use mullvad_api::{proxy::ApiConnectionMode, ApiEndpoint, API};
 
     let api_endpoint = ApiEndpoint::from_env_vars();
     let api_host = format!("api.{}", TEST_CONFIG.mullvad_host);
-    let api_address = format!("{api_host}:443")
-        .to_socket_addrs()
-        .expect("failed to resolve API host")
-        .next()
-        .unwrap();
+
+    let api_host_with_port = format!("{api_host}:443");
+    let api_address = helpers::resolve_hostname_with_retries(api_host_with_port)
+        .await
+        .context("failed to resolve API host")?;
 
     // Override the API endpoint to use the one specified in the test config
     let _ = API.override_init(ApiEndpoint {
@@ -237,7 +240,7 @@ pub fn new_device_client() -> DevicesProxy {
     let api = mullvad_api::Runtime::new(tokio::runtime::Handle::current())
         .expect("failed to create api runtime");
     let rest_handle = api.mullvad_rest_handle(ApiConnectionMode::Direct.into_provider());
-    DevicesProxy::new(rest_handle)
+    Ok(DevicesProxy::new(rest_handle))
 }
 
 pub async fn list_devices_with_retries(
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index fbf447df5d..96c6dc65bb 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -223,6 +223,36 @@ pub async fn ping_sized_with_timeout(
         .map_err(Error::Rpc)
 }
 
+/// Return the first address that `host` resolves to
+pub async fn resolve_hostname_with_retries(
+    host: impl tokio::net::ToSocketAddrs,
+) -> std::io::Result<SocketAddr> {
+    const MAX_ATTEMPTS: usize = 10;
+    const RETRY_DELAY: Duration = Duration::from_secs(5);
+
+    let mut last_error = None;
+    let mut attempt = 0;
+
+    loop {
+        if attempt >= MAX_ATTEMPTS {
+            break Err(last_error.unwrap_or(std::io::Error::other("lookup timed out")));
+        }
+        attempt += 1;
+
+        match tokio::net::lookup_host(&host).await {
+            Ok(mut addrs) => {
+                if let Some(addr) = addrs.next() {
+                    // done
+                    break Ok(addr);
+                }
+            }
+            Err(err) => last_error = Some(err),
+        }
+
+        tokio::time::sleep(RETRY_DELAY).await;
+    }
+}
+
 /// Log in and retry if it fails due to throttling
 pub async fn login_with_retries(
     mullvad_client: &mut MullvadProxyClient,
@@ -501,21 +531,19 @@ pub fn unreachable_wireguard_tunnel() -> talpid_types::net::wireguard::Connectio
 /// # Note
 /// This is independent of the running daemon's environment.
 /// It is solely dependant on the current value of [`TEST_CONFIG`].
-pub fn get_app_env() -> HashMap<String, String> {
+pub async fn get_app_env() -> anyhow::Result<HashMap<String, String>> {
     use mullvad_api::env;
-    use std::net::ToSocketAddrs;
 
     let api_host = format!("api.{}", TEST_CONFIG.mullvad_host);
-    let api_addr = format!("{api_host}:443")
-        .to_socket_addrs()
-        .expect("failed to resolve API host")
-        .next()
-        .unwrap();
+    let api_host_with_port = format!("{api_host}:443");
+    let api_addr = resolve_hostname_with_retries(api_host_with_port)
+        .await
+        .context("failed to resolve API host")?;
 
-    HashMap::from_iter(vec![
+    Ok(HashMap::from_iter(vec![
         (env::API_HOST_VAR.to_string(), api_host),
         (env::API_ADDR_VAR.to_string(), api_addr.to_string()),
-    ])
+    ]))
 }
 
 /// Constrain the daemon to only select the relay selected with `query` when establishing all
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index f3541bcb52..d7056cce76 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -1,4 +1,4 @@
-use anyhow::Context;
+use anyhow::{bail, Context};
 use std::time::Duration;
 
 use mullvad_management_interface::MullvadProxyClient;
@@ -14,10 +14,10 @@ use super::{
 
 /// Install the last stable version of the app and verify that it is running.
 #[test_function(priority = -200)]
-pub async fn test_install_previous_app(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_install_previous_app(_: TestContext, rpc: ServiceClient) -> anyhow::Result<()> {
     // verify that daemon is not already running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::NotRunning {
-        return Err(Error::DaemonRunning);
+        bail!(Error::DaemonRunning);
     }
 
     // install package
@@ -27,13 +27,13 @@ pub async fn test_install_previous_app(_: TestContext, rpc: ServiceClient) -> Re
 
     // verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        return Err(Error::DaemonNotRunning);
+        bail!(Error::DaemonNotRunning);
     }
 
     replace_openvpn_cert(&rpc).await?;
 
     // Override env vars
-    rpc.set_daemon_environment(get_app_env()).await?;
+    rpc.set_daemon_environment(get_app_env().await?).await?;
 
     Ok(())
 }
@@ -45,15 +45,18 @@ pub async fn test_install_previous_app(_: TestContext, rpc: ServiceClient) -> Re
 /// * The installer does not successfully complete.
 /// * The VPN service is not running after the upgrade.
 #[test_function(priority = -190)]
-pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::Result<()> {
     // Verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        return Err(Error::DaemonNotRunning);
+        bail!(Error::DaemonNotRunning);
     }
 
-    super::account::clear_devices(&super::account::new_device_client())
+    let device_client = super::account::new_device_client()
         .await
-        .expect("failed to clear devices");
+        .context("Failed to create device client")?;
+    super::account::clear_devices(&device_client)
+        .await
+        .context("failed to clear devices")?;
 
     // Login to test preservation of device/account
     // TODO: Cannot do this now because overriding the API is impossible for releases
@@ -68,10 +71,10 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<()
 
     rpc.exec("mullvad", ["debug", "block-connection"])
         .await
-        .expect("Failed to set relay location");
+        .context("Failed to set relay location")?;
     rpc.exec("mullvad", ["connect"])
         .await
-        .expect("Failed to begin connecting");
+        .context("Failed to begin connecting")?;
 
     tokio::time::timeout(super::WAIT_FOR_TUNNEL_STATE_TIMEOUT, async {
         // use polling for sake of simplicity
@@ -108,7 +111,7 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<()
 
     // verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        return Err(Error::DaemonNotRunning);
+        bail!(Error::DaemonNotRunning);
     }
 
     // Check if any traffic was observed
@@ -131,7 +134,7 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> Result<()
     let settings = mullvad_client
         .get_settings()
         .await
-        .expect("failed to obtain settings");
+        .context("failed to obtain settings")?;
 
     const EXPECTED_COUNTRY: &str = "xx";
 
@@ -183,9 +186,9 @@ pub async fn test_uninstall_app(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        return Err(Error::DaemonNotRunning);
+        bail!(Error::DaemonNotRunning);
     }
 
     // Login to test preservation of device/account
@@ -193,21 +196,21 @@ pub async fn test_uninstall_app(
     mullvad_client
         .login_account(TEST_CONFIG.account_number.clone())
         .await
-        .expect("login failed");
+        .context("login failed")?;
 
     // save device to verify that uninstalling removes the device
     // we should still be logged in after upgrading
     let uninstalled_device = mullvad_client
         .get_device()
         .await
-        .expect("failed to get device data")
+        .context("failed to get device data")?
         .into_device()
-        .expect("failed to get device")
+        .context("failed to get device")?
         .device
         .id;
 
     log::debug!("Uninstalling app");
-    rpc.uninstall_app(get_app_env()).await?;
+    rpc.uninstall_app(get_app_env().await?).await?;
 
     let app_traces = rpc
         .find_mullvad_app_traces()
@@ -219,11 +222,14 @@ pub async fn test_uninstall_app(
     );
 
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::NotRunning {
-        return Err(Error::DaemonRunning);
+        bail!(Error::DaemonRunning);
     }
 
     // verify that device was removed
-    let devices = super::account::list_devices_with_retries(&super::account::new_device_client())
+    let device_client = super::account::new_device_client()
+        .await
+        .context("Failed to create device client")?;
+    let devices = super::account::list_devices_with_retries(&device_client)
         .await
         .expect("failed to list devices");
 
@@ -239,10 +245,10 @@ pub async fn test_uninstall_app(
 /// Install the app cleanly, failing if the installer doesn't succeed
 /// or if the VPN service is not running afterwards.
 #[test_function(always_run = true, must_succeed = true, priority = -160)]
-pub async fn test_install_new_app(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_install_new_app(_: TestContext, rpc: ServiceClient) -> anyhow::Result<()> {
     // verify that daemon is not already running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::NotRunning {
-        return Err(Error::DaemonRunning);
+        bail!(Error::DaemonRunning);
     }
 
     // install package
@@ -252,7 +258,7 @@ pub async fn test_install_new_app(_: TestContext, rpc: ServiceClient) -> Result<
 
     // verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        return Err(Error::DaemonNotRunning);
+        bail!(Error::DaemonNotRunning);
     }
 
     // Set the log level to trace
@@ -262,7 +268,7 @@ pub async fn test_install_new_app(_: TestContext, rpc: ServiceClient) -> Result<
     replace_openvpn_cert(&rpc).await?;
 
     // Override env vars
-    rpc.set_daemon_environment(get_app_env()).await?;
+    rpc.set_daemon_environment(get_app_env().await?).await?;
 
     Ok(())
 }
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index 8db87e4eae..25811f7d45 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -194,7 +194,7 @@ pub async fn cleanup_after_test(
 /// restart.
 async fn restart_daemon(rpc: ServiceClient) -> anyhow::Result<()> {
     let current_env = rpc.get_daemon_environment().await?;
-    let default_env = get_app_env();
+    let default_env = get_app_env().await?;
     if current_env != default_env {
         log::debug!("Restarting daemon due changed environment variables. Values since last test {current_env:?}");
         rpc.set_daemon_environment(default_env).await?;
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index 7adaf432af..65093a6758 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -134,7 +134,7 @@ async fn test_custom_access_methods_gui(
     _: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
-) -> Result<(), Error> {
+) -> anyhow::Result<()> {
     use mullvad_api::env;
     use mullvad_relay_selector::{RelaySelector, SelectorConfig};
     use talpid_types::net::proxy::CustomProxy;
@@ -155,7 +155,7 @@ async fn test_custom_access_methods_gui(
     // is accomplished by setting the env variable
     // `MULLVAD_API_FORCE_DIRECT=false` and restarting the daemon.
 
-    let mut env = helpers::get_app_env();
+    let mut env = helpers::get_app_env().await?;
     env.insert(env::API_FORCE_DIRECT_VAR.to_string(), "0".to_string());
 
     tokio::time::timeout(