Refactor test cleanup logic

Cleanup is now done BEFORE tests are run and takes care of resetting
the daemon state more thoroughly. The daemon will now always be
installed, logged in and disconnected with all settings reset before
the next test. Tests are therefore not able to depend on the previous
test leaving the test-runner in a certain state and must instead take
care of setting up their own state themselves.

`test_upgrade_app` gets special treatment to be able to run before
the new app version is automatically installed.

Refactor `run_tests.rs`

10 files changed, 481 insertions, 360 deletions

diff --git a/test/test-manager/src/logging.rs b/test/test-manager/src/logging.rs
index 2bc5ea90d1..7484c7833d 100644
--- a/test/test-manager/src/logging.rs
+++ b/test/test-manager/src/logging.rs
@@ -3,6 +3,8 @@ use colored::Colorize;
 use std::sync::{Arc, Mutex};
 use test_rpc::logging::{LogOutput, Output};
 
+use crate::summary;
+
 /// Logger that optionally supports logging records to a buffer
 #[derive(Clone)]
 pub struct Logger {
@@ -109,25 +111,94 @@ impl log::Log for Logger {
     fn flush(&self) {}
 }
 
+/// Encapsulate caught unwound panics, such that we can catch tests that panic and differentiate
+/// them from tests that just fail.
 #[derive(Debug, thiserror::Error)]
-#[error("Test panic: {0}")]
-pub struct PanicMessage(String);
+#[error("Test panic: {}", self.as_string())]
+pub struct Panic(Box<dyn std::any::Any + Send + 'static>);
+
+impl Panic {
+    /// Create a new [`Panic`] from a caught unwound panic.
+    pub fn new(result: Box<dyn std::any::Any + Send + 'static>) -> Self {
+        Self(result)
+    }
+
+    /// Convert this panic to a [`String`] representation.
+    pub fn as_string(&self) -> String {
+        if let Some(result) = self.0.downcast_ref::<String>() {
+            return result.clone();
+        }
+        match self.0.downcast_ref::<&str>() {
+            Some(s) => String::from(*s),
+            None => String::from("unknown message"),
+        }
+    }
+}
 
 pub struct TestOutput {
     pub error_messages: Vec<Output>,
     pub test_name: &'static str,
-    pub result: Result<Result<(), Error>, PanicMessage>,
+    pub result: TestResult,
     pub log_output: Option<LogOutput>,
 }
 
+// Convert this unwieldy return type to a workable `TestResult`.
+// What we are converting from is the acutal return type of the test execution.
+impl From<Result<Result<(), Error>, Panic>> for TestResult {
+    fn from(value: Result<Result<(), Error>, Panic>) -> Self {
+        match value {
+            Ok(Ok(())) => TestResult::Pass,
+            Ok(Err(e)) => TestResult::Fail(e),
+            Err(e) => TestResult::Panic(e),
+        }
+    }
+}
+
+/// Result from a test execution. This may carry information in case the test failed during
+/// execution.
+pub enum TestResult {
+    /// Test passed.
+    Pass,
+    /// Test failed during execution. Contains the source error which caused the test to fail.
+    Fail(Error),
+    /// Test panicked during execution. Contains the caught unwound panic.
+    Panic(Panic),
+}
+
+impl TestResult {
+    /// Returns `true` if test failed or panicked, i.e. when `TestResult` is `Fail` or `Panic`.
+    pub const fn failure(&self) -> bool {
+        matches!(self, TestResult::Fail(_) | TestResult::Panic(_))
+    }
+
+    /// Convert `self` to a [`summary::TestResult`], which is used for creating fancy exports of
+    /// the results for a test run.
+    pub const fn summary(&self) -> summary::TestResult {
+        match self {
+            TestResult::Pass => summary::TestResult::Pass,
+            TestResult::Fail(_) | TestResult::Panic(_) => summary::TestResult::Fail,
+        }
+    }
+
+    /// Consume `self` and convert into a [`Result`] where [`TestResult::Pass`] is mapped to [`Ok`]
+    /// while [`TestResult::Fail`] & [`TestResult::Panic`] is mapped to [`Err`].
+    pub fn anyhow(self) -> anyhow::Result<()> {
+        match self {
+            TestResult::Pass => Ok(()),
+            TestResult::Fail(error) => anyhow::bail!(error),
+            TestResult::Panic(error) => anyhow::bail!(error.to_string()),
+        }
+    }
+}
+
 impl TestOutput {
     pub fn print(&self) {
         match &self.result {
-            Ok(Ok(_)) => {
+            TestResult::Pass => {
                 println!("{}", format!("TEST {} SUCCEEDED!", self.test_name).green());
                 return;
             }
-            Ok(Err(e)) => {
+            TestResult::Fail(e) => {
                 println!(
                     "{}",
                     format!(
@@ -138,13 +209,13 @@ impl TestOutput {
                     .red()
                 );
             }
-            Err(panic_msg) => {
+            TestResult::Panic(panic_msg) => {
                 println!(
                     "{}",
                     format!(
                         "TEST {} PANICKED WITH MESSAGE: {}",
                         self.test_name,
-                        panic_msg.0.bold()
+                        panic_msg.as_string().bold()
                     )
                     .red()
                 );
@@ -191,13 +262,3 @@ impl TestOutput {
         println!("{}", format!("TEST {} END OF OUTPUT", self.test_name).red());
     }
 }
-
-pub fn panic_as_string(error: Box<dyn std::any::Any + Send + 'static>) -> PanicMessage {
-    if let Some(result) = error.downcast_ref::<String>() {
-        return PanicMessage(result.clone());
-    }
-    match error.downcast_ref::<&str>() {
-        Some(s) => PanicMessage(String::from(*s)),
-        None => PanicMessage(String::from("unknown message")),
-    }
-}
diff --git a/test/test-manager/src/main.rs b/test/test-manager/src/main.rs
index b36d94ab7e..9c05289936 100644
--- a/test/test-manager/src/main.rs
+++ b/test/test-manager/src/main.rs
@@ -92,8 +92,9 @@ enum Commands {
         #[arg(long)]
         app_package: String,
 
-        /// App package to upgrade from when running `test_install_previous_app`, can be left empty
-        /// if this test is not ran. Parsed the same way as `--app-package`.
+        /// Given this argument, the `test_upgrade_app` test will run, which installs the previous
+        /// version then upgrades to the version specified in by `--app-package`. If left empty,
+        /// the test will be skipped. Parsed the same way as `--app-package`.
         ///
         /// # Note
         ///
@@ -336,7 +337,8 @@ async fn main() -> Result<()> {
                 instance.wait().await;
             }
             socks.close();
-            result
+            // Propagate any error from the test run if applicable
+            result?.anyhow()
         }
         Commands::FormatTestReports { reports } => {
             summary::print_summary_table(&reports).await;
diff --git a/test/test-manager/src/package.rs b/test/test-manager/src/package.rs
index 34fb30419a..c644d60b72 100644
--- a/test/test-manager/src/package.rs
+++ b/test/test-manager/src/package.rs
@@ -69,7 +69,7 @@ pub fn get_app_manifest(
     })
 }
 
-fn get_version_from_path(app_package_path: &Path) -> Result<String, anyhow::Error> {
+pub fn get_version_from_path(app_package_path: &Path) -> Result<String, anyhow::Error> {
     static VERSION_REGEX: Lazy<Regex> =
         Lazy::new(|| Regex::new(r"\d{4}\.\d+((-beta\d+)?(-dev)?-([0-9a-z])+)?").unwrap());
 
diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 9b6bc7d721..14522d8a84 100644
--- a/test/test-manager/src/run_tests.rs
+++ b/test/test-manager/src/run_tests.rs
@@ -1,14 +1,14 @@
 use crate::{
-    logging::{panic_as_string, TestOutput},
-    mullvad_daemon::{self, MullvadClientArgument},
-    summary::{self, maybe_log_test_result},
+    logging::{Logger, Panic, TestOutput, TestResult},
+    mullvad_daemon::{self, MullvadClientArgument, RpcClientProvider},
+    summary::SummaryLogger,
     tests::{self, config::TEST_CONFIG, get_tests, TestContext},
     vm,
 };
 use anyhow::{Context, Result};
 use futures::FutureExt;
 use std::{future::Future, panic, time::Duration};
-use test_rpc::{logging::Output, mullvad_daemon::MullvadClientVersion, ServiceClient};
+use test_rpc::{logging::Output, ServiceClient};
 
 /// The baud rate of the serial connection between the test manager and the test runner.
 /// There is a known issue with setting a baud rate at all or macOS, and the workaround
@@ -17,14 +17,97 @@ use test_rpc::{logging::Output, mullvad_daemon::MullvadClientVersion, ServiceCli
 /// Keep this constant in sync with `test-runner/src/main.rs`
 const BAUD: u32 = if cfg!(target_os = "macos") { 0 } else { 115200 };
 
+struct TestHandler<'a> {
+    rpc_provider: &'a RpcClientProvider,
+    test_runner_client: &'a ServiceClient,
+    failed_tests: Vec<&'static str>,
+    successful_tests: Vec<&'static str>,
+    summary_logger: Option<SummaryLogger>,
+    print_failed_tests_only: bool,
+    logger: Logger,
+}
+
+impl TestHandler<'_> {
+    /// Run `tests::test_upgrade_app` and register the result
+    async fn run_test<R, F>(
+        &mut self,
+        test: &F,
+        test_name: &'static str,
+        mullvad_client: MullvadClientArgument,
+    ) -> Result<(), anyhow::Error>
+    where
+        F: Fn(super::tests::TestContext, ServiceClient, MullvadClientArgument) -> R,
+        R: Future<Output = anyhow::Result<()>>,
+    {
+        log::info!("Running {test_name}");
+
+        if self.print_failed_tests_only {
+            // Stop live record
+            self.logger.store_records(true);
+        }
+
+        let test_output = run_test_function(
+            self.test_runner_client.clone(),
+            mullvad_client,
+            &test,
+            test_name,
+            TestContext {
+                rpc_provider: self.rpc_provider.clone(),
+            },
+        )
+        .await;
+
+        if self.print_failed_tests_only {
+            // Print results of failed test
+            if test_output.result.failure() {
+                self.logger.print_stored_records();
+            } else {
+                self.logger.flush_records();
+            }
+            self.logger.store_records(false);
+        }
+
+        test_output.print();
+
+        register_test_result(
+            test_output.result,
+            &mut self.failed_tests,
+            test_name,
+            &mut self.successful_tests,
+            self.summary_logger.as_mut(),
+        )
+        .await?;
+
+        Ok(())
+    }
+
+    fn gather_results(self) -> TestResult {
+        log::info!("TESTS THAT SUCCEEDED:");
+        for test in self.successful_tests {
+            log::info!("{test}");
+        }
+
+        log::info!("TESTS THAT FAILED:");
+        for test in &self.failed_tests {
+            log::info!("{test}");
+        }
+
+        if self.failed_tests.is_empty() {
+            TestResult::Pass
+        } else {
+            TestResult::Fail(anyhow::anyhow!("Some tests failed"))
+        }
+    }
+}
+
 pub async fn run(
     config: tests::config::TestConfig,
     instance: &dyn vm::VmInstance,
     test_filters: &[String],
     skip_wait: bool,
     print_failed_tests_only: bool,
-    mut summary_logger: Option<summary::SummaryLogger>,
-) -> Result<()> {
+    summary_logger: Option<SummaryLogger>,
+) -> Result<TestResult> {
     log::trace!("Setting test constants");
     TEST_CONFIG.init(config);
 
@@ -43,11 +126,20 @@ pub async fn run(
 
     log::info!("Running client");
 
-    let client = ServiceClient::new(connection_handle.clone(), runner_transport);
-    let mullvad_client =
-        mullvad_daemon::new_rpc_client(connection_handle, mullvad_daemon_transport);
+    let test_runner_client = ServiceClient::new(connection_handle.clone(), runner_transport);
+    let rpc_provider = mullvad_daemon::new_rpc_client(connection_handle, mullvad_daemon_transport);
 
-    print_os_version(&client).await;
+    print_os_version(&test_runner_client).await;
+
+    let mut test_handler = TestHandler {
+        rpc_provider: &rpc_provider,
+        test_runner_client: &test_runner_client,
+        failed_tests: vec![],
+        successful_tests: vec![],
+        summary_logger,
+        print_failed_tests_only,
+        logger: Logger::get_or_init(),
+    };
 
     let mut tests = get_tests();
 
@@ -68,116 +160,65 @@ pub async fn run(
         });
     }
 
-    let mut final_result = Ok(());
-
-    let test_context = TestContext {
-        rpc_provider: mullvad_client,
+    if TEST_CONFIG.app_package_to_upgrade_from_filename.is_some() {
+        test_handler
+            .run_test(
+                &tests::test_upgrade_app,
+                "test_upgrade_app",
+                MullvadClientArgument::None,
+            )
+            .await?;
+    } else {
+        log::warn!("No previous app to upgrade from, skipping upgrade test");
     };
 
-    let mut successful_tests = vec![];
-    let mut failed_tests = vec![];
-
-    let logger = super::logging::Logger::get_or_init();
-
     for test in tests {
-        let mullvad_client = test_context
-            .rpc_provider
+        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;
+        test_handler
+            .run_test(&test.func, test.name, mullvad_client)
+            .await?;
+    }
 
-        log::info!("Running {}", test.name);
-
-        if print_failed_tests_only {
-            // Stop live record
-            logger.store_records(true);
-        }
-
-        // TODO: Log how long each test took to run.
-        let test_result = run_test(
-            client.clone(),
-            mullvad_client,
-            &test.func,
-            test.name,
-            test_context.clone(),
-        )
-        .await;
-
-        if test.mullvad_client_version == MullvadClientVersion::New {
-            // Try to reset the daemon state if the test failed OR if the test doesn't explicitly
-            // disabled cleanup.
-            if test.cleanup || matches!(test_result.result, Err(_) | Ok(Err(_))) {
-                crate::tests::cleanup_after_test(client.clone(), &test_context.rpc_provider)
-                    .await?;
-            }
-        }
-
-        if print_failed_tests_only {
-            // Print results of failed test
-            if matches!(test_result.result, Err(_) | Ok(Err(_))) {
-                logger.print_stored_records();
-            } else {
-                logger.flush_records();
-            }
-            logger.store_records(false);
-        }
-
-        test_result.print();
-
-        let test_succeeded = matches!(test_result.result, Ok(Ok(_)));
+    let result = test_handler.gather_results();
 
-        maybe_log_test_result(
-            summary_logger.as_mut(),
-            test.name,
-            if test_succeeded {
-                summary::TestResult::Pass
-            } else {
-                summary::TestResult::Fail
-            },
-        )
-        .await
-        .context("Failed to log test result")?;
+    // wait for cleanup
+    drop(test_runner_client);
+    drop(rpc_provider);
+    let _ = tokio::time::timeout(Duration::from_secs(5), completion_handle).await;
 
-        match test_result.result {
-            Err(panic) => {
-                failed_tests.push(test.name);
-                final_result = Err(panic).context("test panicked");
-                if test.must_succeed {
-                    break;
-                }
-            }
-            Ok(Err(failure)) => {
-                failed_tests.push(test.name);
-                final_result = Err(failure).context("test failed");
-                if test.must_succeed {
-                    break;
-                }
-            }
-            Ok(Ok(result)) => {
-                successful_tests.push(test.name);
-                final_result = final_result.and(Ok(result));
-            }
-        }
-    }
+    Ok(result)
+}
 
-    log::info!("TESTS THAT SUCCEEDED:");
-    for test in successful_tests {
-        log::info!("{test}");
-    }
+async fn register_test_result(
+    test_result: TestResult,
+    failed_tests: &mut Vec<&str>,
+    test_name: &'static str,
+    successful_tests: &mut Vec<&str>,
+    summary_logger: Option<&mut SummaryLogger>,
+) -> anyhow::Result<()> {
+    if let Some(logger) = summary_logger {
+        logger
+            .log_test_result(test_name, test_result.summary())
+            .await
+            .context("Failed to log test result")?
+    };
 
-    log::info!("TESTS THAT FAILED:");
-    for test in failed_tests {
-        log::info!("{test}");
+    if test_result.failure() {
+        failed_tests.push(test_name);
+    } else {
+        successful_tests.push(test_name);
     }
 
-    // wait for cleanup
-    drop(client);
-    drop(test_context);
-    let _ = tokio::time::timeout(Duration::from_secs(5), completion_handle).await;
-
-    final_result
+    Ok(())
 }
 
-pub async fn run_test<F, R>(
+pub async fn run_test_function<F, R>(
     runner_rpc: ServiceClient,
     mullvad_rpc: MullvadClientArgument,
     test: &F,
@@ -194,13 +235,15 @@ where
     // assertion being incorrect can not lead to memory unsafety however it could theoretically
     // lead to logic bugs. The problem of forcing the test to be unwind safe is that it causes a
     // large amount of unergonomic design.
-    let result = panic::AssertUnwindSafe(test(test_context, runner_rpc.clone(), mullvad_rpc))
-        .catch_unwind()
-        .await
-        .map_err(panic_as_string);
+    let result: TestResult =
+        panic::AssertUnwindSafe(test(test_context, runner_rpc.clone(), mullvad_rpc))
+            .catch_unwind()
+            .await
+            .map_err(Panic::new)
+            .into();
 
     let mut output = vec![];
-    if matches!(result, Ok(Err(_)) | Err(_)) {
+    if result.failure() {
         let output_after_test = runner_rpc.try_poll_output().await;
         match output_after_test {
             Ok(mut output_after_test) => {
diff --git a/test/test-manager/src/summary.rs b/test/test-manager/src/summary.rs
index 9706e351a1..ab217642d3 100644
--- a/test/test-manager/src/summary.rs
+++ b/test/test-manager/src/summary.rs
@@ -107,19 +107,6 @@ impl SummaryLogger {
     }
 }
 
-/// Convenience function that logs when there's a value, and is a no-op otherwise.
-// y u no trait async fn
-pub async fn maybe_log_test_result(
-    summary_logger: Option<&mut SummaryLogger>,
-    test_name: &str,
-    test_result: TestResult,
-) -> Result<(), Error> {
-    match summary_logger {
-        Some(logger) => logger.log_test_result(test_name, test_result).await,
-        None => Ok(()),
-    }
-}
-
 /// Parsed summary results
 pub struct Summary {
     /// Name of the configuration
diff --git a/test/test-manager/src/tests/account.rs b/test/test-manager/src/tests/account.rs
index d0271a007f..7fe14ae58e 100644
--- a/test/test-manager/src/tests/account.rs
+++ b/test/test-manager/src/tests/account.rs
@@ -9,36 +9,10 @@ use mullvad_types::{
     states::TunnelState,
 };
 use std::time::Duration;
-use talpid_types::net::wireguard;
-use talpid_types::net::wireguard::PublicKey;
+use talpid_types::net::{wireguard, wireguard::PublicKey};
 use test_macro::test_function;
 use test_rpc::ServiceClient;
 
-/// Log in and create a new device for the account.
-#[test_function(always_run = true, must_succeed = true, priority = -100)]
-pub async fn test_login(
-    _: TestContext,
-    _rpc: ServiceClient,
-    mut mullvad_client: MullvadProxyClient,
-) -> anyhow::Result<()> {
-    // Instruct daemon to log in
-    //
-
-    clear_devices(&new_device_client().await?)
-        .await
-        .context("failed to clear devices")?;
-
-    log::info!("Logging in/generating device");
-    login_with_retries(&mut mullvad_client)
-        .await
-        .context("login failed")?;
-
-    // Wait for the relay list to be updated
-    helpers::ensure_updated_relay_list(&mut mullvad_client).await?;
-
-    Ok(())
-}
-
 /// Log out and remove the current device
 /// from the account.
 #[test_function(priority = 100)]
@@ -70,30 +44,33 @@ pub async fn test_too_many_devices(
 
     const MAX_ATTEMPTS: usize = 15;
 
-    for _ in 0..MAX_ATTEMPTS {
-        let pubkey = wireguard::PrivateKey::new_from_random().public_key();
+    let fill_devices = || async {
+        for _ in 0..MAX_ATTEMPTS {
+            let pubkey = wireguard::PrivateKey::new_from_random().public_key();
 
-        match device_client
-            .create(TEST_CONFIG.account_number.clone(), pubkey)
-            .await
-        {
-            Ok(_) => (),
-            Err(mullvad_api::rest::Error::ApiError(_status, ref code))
-                if code == mullvad_api::MAX_DEVICES_REACHED =>
+            match device_client
+                .create(TEST_CONFIG.account_number.clone(), pubkey)
+                .await
             {
-                break;
-            }
-            Err(error) => {
-                log::error!(
-                    "Failed to generate device: {error:?}. Retrying after {} seconds",
-                    THROTTLE_RETRY_DELAY.as_secs()
-                );
-                // Sleep for an overly long time.
-                // TODO: Only sleep for this long if the error is caused by throttling.
-                tokio::time::sleep(THROTTLE_RETRY_DELAY).await;
+                Ok(_) => (),
+                Err(mullvad_api::rest::Error::ApiError(_status, ref code))
+                    if code == mullvad_api::MAX_DEVICES_REACHED =>
+                {
+                    break;
+                }
+                Err(error) => {
+                    log::error!(
+                        "Failed to generate device: {error:?}. Retrying after {} seconds",
+                        THROTTLE_RETRY_DELAY.as_secs()
+                    );
+                    // Sleep for an overly long time.
+                    // TODO: Only sleep for this long if the error is caused by throttling.
+                    tokio::time::sleep(THROTTLE_RETRY_DELAY).await;
+                }
             }
         }
-    }
+    };
+    fill_devices().await;
 
     log::info!("Log in with too many devices");
     let login_result = login_with_retries(&mut mullvad_client).await;
@@ -106,6 +83,9 @@ pub async fn test_too_many_devices(
         "Expected too many devices error, got {login_result:?}"
     );
 
+    mullvad_client.logout_account().await?;
+    fill_devices().await;
+
     // Run UI test
     let ui_result = ui::run_test_env(
         &rpc,
@@ -137,10 +117,7 @@ pub async fn test_revoked_device(
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    log::info!("Logging in/generating device");
-    login_with_retries(&mut mullvad_client)
-        .await
-        .context("login failed")?;
+    mullvad_client.connect_tunnel().await?;
 
     let device_id = mullvad_client
         .get_device()
@@ -151,8 +128,6 @@ pub async fn test_revoked_device(
         .device
         .id;
 
-    helpers::connect_and_wait(&mut mullvad_client).await?;
-
     log::debug!("Removing current device");
 
     let device_client = new_device_client()
@@ -249,8 +224,8 @@ pub async fn test_automatic_wireguard_rotation(
 
     // If key has not yet been updated, listen for changes to it
     if new_key == old_key {
-        // Verify rotation has happened within `ROTATION_TIMEOUT` - if the key hasn't been rotated after
-        // that, the rotation probably won't happen anytime soon.
+        // Verify rotation has happened within `ROTATION_TIMEOUT` - if the key hasn't been rotated
+        // after that, the rotation probably won't happen anytime soon.
         log::info!("Listening for device daemon event");
         let device_event = |daemon_event| match daemon_event {
             DaemonEvent::Device(device_event) => Some(device_event),
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 5dba49665c..db236a822a 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -1,6 +1,13 @@
 use super::{config::TEST_CONFIG, Error, TestContext, WAIT_FOR_TUNNEL_STATE_TIMEOUT};
-use crate::network_monitor::{
-    self, start_packet_monitor, MonitorOptions, MonitorUnexpectedlyStopped, PacketMonitor,
+use crate::{
+    mullvad_daemon::RpcClientProvider,
+    network_monitor::{
+        self, start_packet_monitor, MonitorOptions, MonitorUnexpectedlyStopped, PacketMonitor,
+    },
+    tests::{
+        account::{clear_devices, new_device_client},
+        helpers,
+    },
 };
 use anyhow::{anyhow, bail, ensure, Context};
 use futures::StreamExt;
@@ -27,7 +34,9 @@ use std::{
     time::Duration,
 };
 use talpid_types::net::wireguard::{PeerConfig, PrivateKey, TunnelConfig};
-use test_rpc::{meta::Os, package::Package, AmIMullvad, ServiceClient, SpawnOpts};
+use test_rpc::{
+    meta::Os, mullvad_daemon::ServiceStatus, package::Package, AmIMullvad, ServiceClient, SpawnOpts,
+};
 use tokio::time::sleep;
 
 pub const THROTTLE_RETRY_DELAY: Duration = Duration::from_secs(120);
@@ -54,10 +63,64 @@ macro_rules! assert_tunnel_state {
     }};
 }
 
-pub fn get_package_desc(name: &str) -> Result<Package, Error> {
-    Ok(Package {
+/// Install the app cleanly, failing if the installer doesn't succeed
+/// or if the VPN service is not running afterwards.
+pub async fn install_app(
+    rpc: &ServiceClient,
+    app_filename: &str,
+    rpc_provider: &RpcClientProvider,
+) -> anyhow::Result<MullvadProxyClient> {
+    // install package
+    log::info!("Installing app '{}'", app_filename);
+    rpc.install_app(get_package_desc(app_filename)).await?;
+
+    // verify that daemon is running
+    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
+        bail!(Error::DaemonNotRunning);
+    }
+
+    // Set the log level to trace
+    rpc.set_daemon_log_level(test_rpc::mullvad_daemon::Verbosity::Trace)
+        .await?;
+
+    replace_openvpn_certificate(rpc).await?;
+
+    // Override env vars
+    rpc.set_daemon_environment(get_app_env().await?).await?;
+
+    // Wait for the relay list to be updated
+    let mut mullvad_client = rpc_provider.new_client().await;
+    helpers::ensure_updated_relay_list(&mut mullvad_client)
+        .await
+        .context("Failed to update relay list")?;
+    Ok(mullvad_client)
+}
+
+/// Replace the OpenVPN CA certificate which is currently used by the installed Mullvad App.
+/// This needs to be invoked after reach (re)installation to use the custom OpenVPN certificate.
+async fn replace_openvpn_certificate(rpc: &ServiceClient) -> Result<(), Error> {
+    const DEST_CERT_FILENAME: &str = "ca.crt";
+
+    let dest_dir = match TEST_CONFIG.os {
+        Os::Windows => "C:\\Program Files\\Mullvad VPN\\resources",
+        Os::Linux => "/opt/Mullvad VPN/resources",
+        Os::Macos => "/Applications/Mullvad VPN.app/Contents/Resources",
+    };
+
+    let dest = Path::new(dest_dir)
+        .join(DEST_CERT_FILENAME)
+        .as_os_str()
+        .to_string_lossy()
+        .into_owned();
+    rpc.write_file(dest, TEST_CONFIG.openvpn_certificate.to_vec())
+        .await
+        .map_err(Error::Rpc)
+}
+
+pub fn get_package_desc(name: &str) -> Package {
+    Package {
         path: Path::new(&TEST_CONFIG.artifacts_dir).join(name),
-    })
+    }
 }
 
 /// Reboot the guest virtual machine.
@@ -334,15 +397,28 @@ pub async fn login_with_retries(
 /// This will first check whether we are logged in. If not, it will also try to login
 /// on your behalf. If this function returns without any errors, we are logged in to a valid
 /// account.
-pub async fn ensure_logged_in(
-    mullvad_client: &mut MullvadProxyClient,
-) -> Result<(), mullvad_management_interface::Error> {
-    if mullvad_client.get_device().await?.is_logged_in() {
+pub async fn ensure_logged_in(mullvad_client: &mut MullvadProxyClient) -> anyhow::Result<()> {
+    if !matches!(
+        mullvad_client.update_device().await,
+        Err(mullvad_management_interface::Error::DeviceNotFound)
+    ) && mullvad_client.get_device().await?.is_logged_in()
+    {
         return Ok(());
     }
     log::info!("Current device not logged in. Clearing devices and logging in.");
     // We are apparently not logged in already.. Try to log in.
-    login_with_retries(mullvad_client).await
+    clear_devices(
+        &new_device_client()
+            .await
+            .context("Failed to create device client")?,
+    )
+    .await
+    .context("failed to clear devices")?;
+
+    login_with_retries(mullvad_client)
+        .await
+        .context("Failed to log in")?;
+    Ok(())
 }
 
 /// Try to connect to a Mullvad Tunnel.
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index b1688d27ac..2f787f6299 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -1,56 +1,44 @@
-use anyhow::{bail, Context};
-use std::{path::Path, time::Duration};
+use anyhow::{bail, ensure, Context};
+use std::time::Duration;
 
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_types::{constraints::Constraint, relay_constraints};
 use test_macro::test_function;
-use test_rpc::{meta::Os, mullvad_daemon::ServiceStatus, ServiceClient};
+use test_rpc::{mullvad_daemon::ServiceStatus, ServiceClient};
+
+use crate::{mullvad_daemon::MullvadClientArgument, tests::helpers};
 
 use super::{
     config::TEST_CONFIG,
-    helpers::{connect_and_wait, get_app_env, get_package_desc, wait_for_tunnel_state, Pinger},
+    helpers::{
+        connect_and_wait, get_app_env, get_package_desc, install_app, wait_for_tunnel_state, Pinger,
+    },
     Error, TestContext,
 };
 
-/// 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) -> anyhow::Result<()> {
-    // verify that daemon is not already running
-    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::NotRunning {
-        bail!(Error::DaemonRunning);
-    }
-
-    // install package
-    log::debug!("Installing old app");
-    rpc.install_app(get_package_desc(
-        TEST_CONFIG
-            .app_package_to_upgrade_from_filename
-            .as_ref()
-            .context("Missing previous app version")?,
-    )?)
-    .await?;
-    log::debug!("Replacing the OpenVPN CA certificate");
-    replace_openvpn_certificate(&rpc).await?;
-
-    // verify that daemon is running
-    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        bail!(Error::DaemonNotRunning);
-    }
-
-    // Override env vars
-    rpc.set_daemon_environment(get_app_env().await?).await?;
-
-    Ok(())
-}
-
 /// Upgrade to the "version under test". This test fails if:
 ///
 /// * Leaks (TCP/UDP/ICMP) to a single public IP address are successfully produced during the
 ///   upgrade.
 /// * 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) -> anyhow::Result<()> {
+pub async fn test_upgrade_app(
+    ctx: TestContext,
+    rpc: ServiceClient,
+    _mullvad_client: MullvadClientArgument,
+) -> anyhow::Result<()> {
+    // Install the older version of the app and verify that it is running.
+    install_app(
+        &rpc,
+        TEST_CONFIG
+            .app_package_to_upgrade_from_filename
+            .as_ref()
+            .unwrap(),
+        &ctx.rpc_provider,
+    )
+    .await
+    .context("Failed to install previous app version")?;
+
     // Verify that daemon is running
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
         bail!(Error::DaemonNotRunning);
@@ -68,7 +56,7 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
     // mullvad_client
     //    .login_account(TEST_CONFIG.account_number.clone())
     //    .await
-    //    .expect("login failed");
+    //    .context("login failed")?;
 
     // Start blocking
     //
@@ -107,8 +95,9 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
     let pinger = Pinger::start(&rpc).await;
 
     // install new package
+
     log::debug!("Installing new app");
-    rpc.install_app(get_package_desc(&TEST_CONFIG.app_package_filename)?)
+    rpc.install_app(helpers::get_package_desc(&TEST_CONFIG.app_package_filename))
         .await?;
 
     // Give it some time to start
@@ -118,14 +107,12 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
     if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
         bail!(Error::DaemonNotRunning);
     }
-
     // Check if any traffic was observed
     //
     let guest_ip = pinger.guest_ip;
-    let monitor_result = pinger.stop().await.unwrap();
-    assert_eq!(
-        monitor_result.packets.len(),
-        0,
+    let monitor_result = pinger.stop().await.context("Failed to stop pinger")?;
+    ensure!(
+        monitor_result.packets.is_empty(),
         "observed unexpected packets from {guest_ip}"
     );
 
@@ -154,7 +141,7 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
         _ => false,
     };
 
-    assert!(
+    ensure!(
         relay_location_was_preserved,
         "relay location was not preserved after upgrade. new settings: {:?}",
         settings,
@@ -163,13 +150,12 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
     // check if account history was preserved
     // TODO: Cannot check account history because overriding the API is impossible for releases
     // let history = mullvad_client
-    // .get_account_history(())
-    // .await
-    // .expect("failed to obtain account history");
-    // assert_eq!(
-    // history.into_inner().token,
-    // Some(TEST_CONFIG.account_number.clone()),
-    // "lost account history"
+    //     .get_account_history()
+    //     .await
+    //     .context("failed to obtain account history")?;
+    // ensure!(
+    //     history.into_inner().token == Some(TEST_CONFIG.account_number.clone()),
+    //     "lost account history"
     // );
 
     Ok(())
@@ -186,23 +172,12 @@ pub async fn test_upgrade_app(ctx: TestContext, rpc: ServiceClient) -> anyhow::R
 /// Files due to Electron, temporary files, registry
 /// values/keys, and device drivers are not guaranteed
 /// to be deleted.
-#[test_function(priority = -170, cleanup = false)]
+#[test_function(priority = -160)]
 pub async fn test_uninstall_app(
-    _: TestContext,
+    _ctx: TestContext,
     rpc: ServiceClient,
     mut mullvad_client: MullvadProxyClient,
 ) -> anyhow::Result<()> {
-    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        bail!(Error::DaemonNotRunning);
-    }
-
-    // Login to test preservation of device/account
-    // TODO: Remove once we can login before upgrade above
-    mullvad_client
-        .login_account(TEST_CONFIG.account_number.clone())
-        .await
-        .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
@@ -247,37 +222,6 @@ pub async fn test_uninstall_app(
     Ok(())
 }
 
-/// 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) -> anyhow::Result<()> {
-    // verify that daemon is not already running
-    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::NotRunning {
-        bail!(Error::DaemonRunning);
-    }
-
-    // install package
-    log::debug!("Installing new app");
-    rpc.install_app(get_package_desc(&TEST_CONFIG.app_package_filename)?)
-        .await?;
-
-    // verify that daemon is running
-    if rpc.mullvad_daemon_get_status().await? != ServiceStatus::Running {
-        bail!(Error::DaemonNotRunning);
-    }
-
-    // Set the log level to trace
-    rpc.set_daemon_log_level(test_rpc::mullvad_daemon::Verbosity::Trace)
-        .await?;
-
-    replace_openvpn_certificate(&rpc).await?;
-
-    // Override env vars
-    rpc.set_daemon_environment(get_app_env().await?).await?;
-
-    Ok(())
-}
-
 /// Install the multiple times starting from a connected state with auto-connect
 /// disabled, failing if the app starts in a disconnected state.
 ///
@@ -306,14 +250,14 @@ pub async fn test_installation_idempotency(
 
     // Check for traffic leaks during the installation processes.
     //
-    // Start continously pinging while monitoring the network traffic. No
+    // Start continuously pinging while monitoring the network traffic. No
     // traffic should be observed going outside of the tunnel during either
     // installation process.
     let pinger = Pinger::start(&rpc).await;
     for _ in 0..2 {
         // Install the app
         log::info!("Installing new app");
-        let app_package = get_package_desc(&TEST_CONFIG.app_package_filename)?;
+        let app_package = get_package_desc(&TEST_CONFIG.app_package_filename);
         rpc.install_app(app_package).await?;
         log::info!("App was successfully installed!");
 
@@ -347,24 +291,3 @@ pub async fn test_installation_idempotency(
 
     Ok(())
 }
-
-/// Replace the OpenVPN CA certificate which is currently used by the installed Mullvad App.
-/// This needs to be invoked after reach (re)installation to use the custom OpenVPN certificate.
-async fn replace_openvpn_certificate(rpc: &ServiceClient) -> Result<(), Error> {
-    const DEST_CERT_FILENAME: &str = "ca.crt";
-
-    let dest_dir = match TEST_CONFIG.os {
-        Os::Windows => "C:\\Program Files\\Mullvad VPN\\resources",
-        Os::Linux => "/opt/Mullvad VPN/resources",
-        Os::Macos => "/Applications/Mullvad VPN.app/Contents/Resources",
-    };
-
-    let dest = Path::new(dest_dir)
-        .join(DEST_CERT_FILENAME)
-        .as_os_str()
-        .to_string_lossy()
-        .into_owned();
-    rpc.write_file(dest, TEST_CONFIG.openvpn_certificate.to_vec())
-        .await
-        .map_err(Error::Rpc)
-}
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index f68d25ffa3..74502b00b0 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -22,8 +22,12 @@ use std::time::Duration;
 
 use crate::{
     mullvad_daemon::{MullvadClientArgument, RpcClientProvider},
-    tests::helpers::get_app_env,
+    package::get_version_from_path,
 };
+use config::TEST_CONFIG;
+use helpers::{get_app_env, install_app};
+pub use install::test_upgrade_app;
+use mullvad_management_interface::MullvadProxyClient;
 use test_rpc::ServiceClient;
 
 const WAIT_FOR_TUNNEL_STATE_TIMEOUT: Duration = Duration::from_secs(40);
@@ -74,34 +78,93 @@ pub fn get_tests() -> Vec<&'static TestMetadata> {
     tests
 }
 
-/// Restore settings to the defaults.
-pub async fn cleanup_after_test(
-    rpc: ServiceClient,
+/// Make sure the daemon is installed and logged in and restore settings to the defaults.
+pub async fn prepare_daemon(
+    rpc: &ServiceClient,
     rpc_provider: &RpcClientProvider,
 ) -> anyhow::Result<()> {
-    log::debug!("Resetting daemon settings after test");
-    // Check if daemon should be restarted.
-    restart_daemon(rpc).await?;
-    let mut mullvad_client = rpc_provider.new_client().await;
+    // Check if daemon should be restarted
+    let mut mullvad_client = ensure_daemon_version(rpc, rpc_provider)
+        .await
+        .context("Failed to restart daemon")?;
+
+    log::debug!("Resetting daemon settings before test");
     mullvad_client
         .reset_settings()
         .await
         .context("Failed to reset settings")?;
-    helpers::disconnect_and_wait(&mut mullvad_client).await?;
+    helpers::disconnect_and_wait(&mut mullvad_client)
+        .await
+        .context("Failed to disconnect daemon after test")?;
+    helpers::ensure_logged_in(&mut mullvad_client).await?;
+
     Ok(())
 }
 
+/// Reset the daemons environment.
+///
+/// Will and restart or reinstall it if necessary.
+async fn ensure_daemon_version(
+    rpc: &ServiceClient,
+    rpc_provider: &RpcClientProvider,
+) -> anyhow::Result<MullvadProxyClient> {
+    let app_package_filename = &TEST_CONFIG.app_package_filename;
+
+    let mullvad_client = if correct_daemon_version_is_running(rpc_provider.new_client().await).await
+    {
+        ensure_daemon_environment(rpc)
+            .await
+            .context("Failed to reset daemon environment")?;
+        rpc_provider.new_client().await
+    } else {
+        // NOTE: Reinstalling the app resets the daemon environment
+        install_app(rpc, app_package_filename, rpc_provider)
+            .await
+            .with_context(|| format!("Failed to install app '{app_package_filename}'"))?
+    };
+    Ok(mullvad_client)
+}
+
 /// Conditionally restart the running daemon
 ///
 /// If the daemon was started with non-standard environment variables, subsequent tests may break
 /// due to assuming a default configuration. In that case, reset the environment variables and
 /// restart.
-async fn restart_daemon(rpc: ServiceClient) -> anyhow::Result<()> {
-    let current_env = rpc.get_daemon_environment().await?;
-    let default_env = get_app_env().await?;
+pub async fn ensure_daemon_environment(rpc: &ServiceClient) -> Result<(), anyhow::Error> {
+    let current_env = rpc
+        .get_daemon_environment()
+        .await
+        .context("Failed to get daemon env variables")?;
+    let default_env = get_app_env()
+        .await
+        .context("Failed to get daemon default env variables")?;
     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?;
-    }
+        rpc.set_daemon_environment(default_env)
+            .await
+            .context("Failed to restart daemon")?;
+    };
     Ok(())
 }
+
+/// Checks if daemon is installed with the version specified by `TEST_CONFIG.app_package_filename`
+async fn correct_daemon_version_is_running(mut mullvad_client: MullvadProxyClient) -> bool {
+    let app_package_filename = &TEST_CONFIG.app_package_filename;
+    let expected_version = get_version_from_path(std::path::Path::new(app_package_filename))
+        .unwrap_or_else(|_| panic!("Invalid app version: {app_package_filename}"));
+
+    use mullvad_management_interface::Error::*;
+    match mullvad_client.get_current_version().await {
+        // Failing to reach the daemon is a sign that it is not installed
+        Err(Rpc(..)) => {
+            log::debug!("Could not reach active daemon before test, it is not running");
+            false
+        }
+        Err(e) => panic!("Failed to get app version: {e}"),
+        Ok(version) if version == expected_version => true,
+        _ => {
+            log::debug!("Daemon version mismatch");
+            false
+        }
+    }
+}
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index f2b266a4e6..e9008a5baf 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -1,8 +1,4 @@
-use super::{
-    config::TEST_CONFIG,
-    helpers::{self, ensure_logged_in},
-    Error, TestContext,
-};
+use super::{config::TEST_CONFIG, helpers, Error, TestContext};
 use mullvad_management_interface::MullvadProxyClient;
 use mullvad_relay_selector::query::builder::RelayQueryBuilder;
 use std::{
@@ -124,7 +120,12 @@ pub async fn test_ui_tunnel_settings(
 
 /// Test whether logging in and logging out work in the GUI
 #[test_function(priority = 500)]
-pub async fn test_ui_login(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_ui_login(
+    _: TestContext,
+    rpc: ServiceClient,
+    mut mullvad_client: MullvadProxyClient,
+) -> Result<(), Error> {
+    mullvad_client.logout_account().await?;
     let ui_result = run_test_env(
         &rpc,
         &["login.spec"],
@@ -229,13 +230,6 @@ async fn test_custom_bridge_gui(
     // See `gui/test/e2e/installed/state-dependent/custom-bridge.spec.ts`
     // for details. The setup should be the same as in
     // `test_manager::tests::access_methods::test_shadowsocks`.
-    //
-    // # Note
-    // The test requires the app to already be logged in.
-
-    ensure_logged_in(&mut mullvad_client)
-        .await
-        .expect("ensure_logged_in failed");
 
     let gui_test = "custom-bridge.spec";
     let relay_list = mullvad_client.get_relay_locations().await.unwrap();
@@ -276,9 +270,6 @@ async fn test_custom_bridge_gui(
 }
 
 /// Test settings import / IP overrides in the GUI
-///
-/// # Note
-/// This test expects the daemon to be logged in
 #[test_function]
 pub async fn test_import_settings_ui(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
     let ui_result = run_test(&rpc, &["settings-import.spec"]).await?;