Merge branch 'be-able-to-whitelist-relays-in-an-end-to-end-test-run-des-1052'

1 files changed, 13 insertions, 15 deletions

diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 88577fb3b3..54e563f704 100644
--- a/test/test-manager/src/run_tests.rs
+++ b/test/test-manager/src/run_tests.rs
@@ -1,12 +1,13 @@
 use crate::{
     logging::{Logger, Panic, TestOutput, TestResult},
-    mullvad_daemon::{self, MullvadClientArgument, RpcClientProvider},
+    mullvad_daemon::{self, RpcClientProvider},
     summary::SummaryLogger,
     tests::{self, config::TEST_CONFIG, TestContext, TestMetadata},
     vm,
 };
 use anyhow::{Context, Result};
 use futures::FutureExt;
+use mullvad_management_interface::MullvadProxyClient;
 use std::{future::Future, panic, time::Duration};
 use test_rpc::{logging::Output, ServiceClient};
 
@@ -33,10 +34,10 @@ impl TestHandler<'_> {
         &mut self,
         test: &F,
         test_name: &'static str,
-        mullvad_client: MullvadClientArgument,
+        mullvad_client: Option<MullvadProxyClient>,
     ) -> Result<(), anyhow::Error>
     where
-        F: Fn(super::tests::TestContext, ServiceClient, MullvadClientArgument) -> R,
+        F: Fn(super::tests::TestContext, ServiceClient, Option<MullvadProxyClient>) -> R,
         R: Future<Output = anyhow::Result<()>>,
     {
         log::info!("Running {test_name}");
@@ -146,26 +147,23 @@ pub async fn run(
     // expected, and to allow for skipping tests on arbitrary conditions.
     if TEST_CONFIG.app_package_to_upgrade_from_filename.is_some() {
         test_handler
-            .run_test(
-                &tests::test_upgrade_app,
-                "test_upgrade_app",
-                MullvadClientArgument::None,
-            )
+            .run_test(&tests::test_upgrade_app, "test_upgrade_app", None)
             .await?;
     } else {
         log::warn!("No previous app to upgrade from, skipping upgrade test");
     };
 
     for test in tests {
-        tests::prepare_daemon(&test_runner_client, &rpc_provider)
+        let mut mullvad_client = 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;
+        tests::set_test_location(&mut mullvad_client, &test)
+            .await
+            .context("Failed to create custom list from test locations")?;
+
         test_handler
-            .run_test(&test.func, test.name, mullvad_client)
+            .run_test(&test.func, test.name, Some(mullvad_client))
             .await?;
     }
 
@@ -204,13 +202,13 @@ async fn register_test_result(
 
 pub async fn run_test_function<F, R>(
     runner_rpc: ServiceClient,
-    mullvad_rpc: MullvadClientArgument,
+    mullvad_rpc: Option<MullvadProxyClient>,
     test: &F,
     test_name: &'static str,
     test_context: super::tests::TestContext,
 ) -> TestOutput
 where
-    F: Fn(super::tests::TestContext, ServiceClient, MullvadClientArgument) -> R,
+    F: Fn(super::tests::TestContext, ServiceClient, Option<MullvadProxyClient>) -> R,
     R: Future<Output = anyhow::Result<()>>,
 {
     let _flushed = runner_rpc.try_poll_output().await;