Remove mullvad client version from macro

Replace `MullvadClientArgument` with `Option`

Small refactor

9 files changed, 47 insertions, 147 deletions

diff --git a/test/test-manager/src/mullvad_daemon.rs b/test/test-manager/src/mullvad_daemon.rs
index b039813470..8da149d329 100644
--- a/test/test-manager/src/mullvad_daemon.rs
+++ b/test/test-manager/src/mullvad_daemon.rs
@@ -4,10 +4,7 @@ use std::{io, time::Duration};
 use futures::{channel::mpsc, future::BoxFuture, pin_mut, FutureExt, SinkExt, StreamExt};
 use hyper_util::rt::TokioIo;
 use mullvad_management_interface::{ManagementServiceClient, MullvadProxyClient};
-use test_rpc::{
-    mullvad_daemon::MullvadClientVersion,
-    transport::{ConnectionHandle, GrpcForwarder},
-};
+use test_rpc::transport::{ConnectionHandle, GrpcForwarder};
 use tokio::io::{AsyncReadExt, AsyncWriteExt, DuplexStream};
 use tokio_util::codec::{Decoder, LengthDelimitedCodec};
 use tower::Service;
@@ -52,20 +49,7 @@ pub struct RpcClientProvider {
     service: DummyService,
 }
 
-pub enum MullvadClientArgument {
-    WithClient(MullvadProxyClient),
-    None,
-}
-
 impl RpcClientProvider {
-    /// Whether a [test case](test_macro::test_function) needs a [`MullvadProxyClient`].
-    pub async fn mullvad_client(&self, client_type: MullvadClientVersion) -> MullvadClientArgument {
-        match client_type {
-            MullvadClientVersion::New => MullvadClientArgument::WithClient(self.new_client().await),
-            MullvadClientVersion::None => MullvadClientArgument::None,
-        }
-    }
-
     pub async fn new_client(&self) -> MullvadProxyClient {
         // FIXME: Ugly workaround to ensure that we don't receive stuff from a
         // previous RPC session.
diff --git a/test/test-manager/src/run_tests.rs b/test/test-manager/src/run_tests.rs
index 3d2f1702e0..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,11 +147,7 @@ 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");
@@ -166,11 +163,7 @@ pub async fn run(
             .context("Failed to create custom list from test locations")?;
 
         test_handler
-            .run_test(
-                &test.func,
-                test.name,
-                MullvadClientArgument::WithClient(mullvad_client),
-            )
+            .run_test(&test.func, test.name, Some(mullvad_client))
             .await?;
     }
 
@@ -209,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;
diff --git a/test/test-manager/src/tests/install.rs b/test/test-manager/src/tests/install.rs
index 936f07d65d..d640b26089 100644
--- a/test/test-manager/src/tests/install.rs
+++ b/test/test-manager/src/tests/install.rs
@@ -6,7 +6,7 @@ use mullvad_types::{constraints::Constraint, relay_constraints};
 use test_macro::test_function;
 use test_rpc::{mullvad_daemon::ServiceStatus, ServiceClient};
 
-use crate::{mullvad_daemon::MullvadClientArgument, tests::helpers};
+use crate::tests::helpers;
 
 use super::{
     config::TEST_CONFIG,
@@ -25,7 +25,7 @@ use super::{
 pub async fn test_upgrade_app(
     ctx: TestContext,
     rpc: ServiceClient,
-    _mullvad_client: MullvadClientArgument,
+    _mullvad_client: Option<MullvadProxyClient>,
 ) -> anyhow::Result<()> {
     // Install the older version of the app and verify that it is running.
     let old_version = TEST_CONFIG
diff --git a/test/test-manager/src/tests/mod.rs b/test/test-manager/src/tests/mod.rs
index b369b28067..e2d50d0889 100644
--- a/test/test-manager/src/tests/mod.rs
+++ b/test/test-manager/src/tests/mod.rs
@@ -23,10 +23,7 @@ use anyhow::Context;
 use futures::future::BoxFuture;
 use std::time::Duration;
 
-use crate::{
-    mullvad_daemon::{MullvadClientArgument, RpcClientProvider},
-    package::get_version_from_path,
-};
+use crate::{mullvad_daemon::RpcClientProvider, package::get_version_from_path};
 use config::TEST_CONFIG;
 use helpers::{find_custom_list, get_app_env, install_app, set_location};
 pub use install::test_upgrade_app;
@@ -40,8 +37,11 @@ pub struct TestContext {
     pub rpc_provider: RpcClientProvider,
 }
 
-pub type TestWrapperFunction =
-    fn(TestContext, ServiceClient, MullvadClientArgument) -> BoxFuture<'static, anyhow::Result<()>>;
+pub type TestWrapperFunction = fn(
+    TestContext,
+    ServiceClient,
+    Option<MullvadProxyClient>,
+) -> BoxFuture<'static, anyhow::Result<()>>;
 
 #[derive(thiserror::Error, Debug)]
 pub enum Error {
diff --git a/test/test-manager/src/tests/split_tunnel.rs b/test/test-manager/src/tests/split_tunnel.rs
index 435552a460..98cbed7951 100644
--- a/test/test-manager/src/tests/split_tunnel.rs
+++ b/test/test-manager/src/tests/split_tunnel.rs
@@ -86,7 +86,11 @@ pub async fn test_split_tunnel(
 /// - A split process should never push traffic through the tunnel.
 /// - Splitting/unsplitting should work regardless if process is running.
 #[test_function(target_os = "macos")]
-pub async fn test_split_tunnel_ui(_ctx: TestContext, rpc: ServiceClient) -> anyhow::Result<()> {
+pub async fn test_split_tunnel_ui(
+    _ctx: TestContext,
+    rpc: ServiceClient,
+    _: MullvadProxyClient,
+) -> anyhow::Result<()> {
     // Skip test on macOS 12, since the feature is unsupported
     if is_macos_12_or_lower(&rpc).await? {
         return Ok(());
diff --git a/test/test-manager/src/tests/test_metadata.rs b/test/test-manager/src/tests/test_metadata.rs
index 961db4510a..31a55498f9 100644
--- a/test/test-manager/src/tests/test_metadata.rs
+++ b/test/test-manager/src/tests/test_metadata.rs
@@ -1,11 +1,10 @@
 use super::TestWrapperFunction;
-use test_rpc::{meta::Os, mullvad_daemon::MullvadClientVersion};
+use test_rpc::meta::Os;
 
 #[derive(Clone, Debug)]
 pub struct TestMetadata {
     pub name: &'static str,
     pub targets: &'static [Os],
-    pub mullvad_client_version: MullvadClientVersion,
     pub func: TestWrapperFunction,
     /// Priority order of the tests, unless specific tests are given as the `TEST_FILTERS` argument
     pub priority: Option<i32>,
diff --git a/test/test-manager/src/tests/ui.rs b/test/test-manager/src/tests/ui.rs
index 9491308eb8..088fd1e55d 100644
--- a/test/test-manager/src/tests/ui.rs
+++ b/test/test-manager/src/tests/ui.rs
@@ -265,7 +265,11 @@ async fn test_custom_bridge_gui(
 
 /// Test settings import / IP overrides in the GUI
 #[test_function]
-pub async fn test_import_settings_ui(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_import_settings_ui(
+    _: TestContext,
+    rpc: ServiceClient,
+    _: MullvadProxyClient,
+) -> Result<(), Error> {
     let ui_result = run_test(&rpc, &["settings-import.spec"]).await?;
     assert!(ui_result.success());
     Ok(())
@@ -273,7 +277,11 @@ pub async fn test_import_settings_ui(_: TestContext, rpc: ServiceClient) -> Resu
 
 /// Test obfuscation settings in the GUI
 #[test_function]
-pub async fn test_obfuscation_settings_ui(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_obfuscation_settings_ui(
+    _: TestContext,
+    rpc: ServiceClient,
+    _: MullvadProxyClient,
+) -> Result<(), Error> {
     let ui_result = run_test(&rpc, &["obfuscation.spec"]).await?;
     assert!(ui_result.success());
     Ok(())
@@ -281,7 +289,11 @@ pub async fn test_obfuscation_settings_ui(_: TestContext, rpc: ServiceClient) ->
 
 /// Test settings in the GUI
 #[test_function]
-pub async fn test_settings_ui(_: TestContext, rpc: ServiceClient) -> Result<(), Error> {
+pub async fn test_settings_ui(
+    _: TestContext,
+    rpc: ServiceClient,
+    _: MullvadProxyClient,
+) -> Result<(), Error> {
     let ui_result = run_test(&rpc, &["settings.spec"]).await?;
     assert!(ui_result.success());
     Ok(())
diff --git a/test/test-manager/test_macro/src/lib.rs b/test/test-manager/test_macro/src/lib.rs
index e16968bd9d..048bb1975e 100644
--- a/test/test-manager/test_macro/src/lib.rs
+++ b/test/test-manager/test_macro/src/lib.rs
@@ -91,11 +91,9 @@ fn parse_marked_test_function(
     function: &syn::ItemFn,
 ) -> Result<TestFunction> {
     let macro_parameters = get_test_macro_parameters(attributes)?;
-    let function_parameters = get_test_function_parameters(&function.sig.inputs)?;
 
     Ok(TestFunction {
         name: function.sig.ident.clone(),
-        function_parameters,
         macro_parameters,
     })
 }
@@ -153,34 +151,15 @@ fn create_test(test_function: TestFunction) -> proc_macro2::TokenStream {
         .collect();
 
     let func_name = test_function.name;
-    let function_mullvad_version = test_function.function_parameters.mullvad_client.version();
-    let wrapper_closure = match test_function.function_parameters.mullvad_client {
-        MullvadClient::New { .. } => {
-            quote! {
-                |test_context: crate::tests::TestContext,
-                rpc: test_rpc::ServiceClient,
-                mullvad_client: crate::mullvad_daemon::MullvadClientArgument|
-                {
-                    let mullvad_client = match mullvad_client {
-                        crate::mullvad_daemon::MullvadClientArgument::WithClient(client) => client,
-                        crate::mullvad_daemon::MullvadClientArgument::None => unreachable!("invalid mullvad client")
-                    };
-                    Box::pin(async move {
-                        #func_name(test_context, rpc, mullvad_client).await.map_err(Into::into)
-                    })
-                }
-            }
-        }
-        MullvadClient::None { .. } => {
-            quote! {
-                |test_context: crate::tests::TestContext,
-                rpc: test_rpc::ServiceClient,
-                _mullvad_client: crate::mullvad_daemon::MullvadClientArgument| {
-                    Box::pin(async move {
-                        #func_name(test_context, rpc).await.map_err(Into::into)
-                    })
-                }
-            }
+    let wrapper_closure = quote! {
+        |test_context: crate::tests::TestContext,
+        rpc: test_rpc::ServiceClient,
+        mullvad_client: Option<MullvadProxyClient>|
+        {
+            let mullvad_client = mullvad_client.expect("Test functions defined using the macro should be given a mullvad client");
+            Box::pin(async move {
+                #func_name(test_context, rpc, mullvad_client).await.map_err(Into::into)
+            })
         }
     };
 
@@ -188,7 +167,6 @@ fn create_test(test_function: TestFunction) -> proc_macro2::TokenStream {
         inventory::submit!(crate::tests::test_metadata::TestMetadata {
             name: stringify!(#func_name),
             targets: &[#targets],
-            mullvad_client_version: #function_mullvad_version,
             func: #wrapper_closure,
             priority: #test_function_priority,
             location: None,
@@ -198,7 +176,6 @@ fn create_test(test_function: TestFunction) -> proc_macro2::TokenStream {
 
 struct TestFunction {
     name: syn::Ident,
-    function_parameters: FunctionParameters,
     macro_parameters: MacroParameters,
 }
 
@@ -206,66 +183,3 @@ struct MacroParameters {
     priority: Option<i32>,
     targets: Vec<Os>,
 }
-
-enum MullvadClient {
-    None {
-        mullvad_client_version: proc_macro2::TokenStream,
-    },
-    New {
-        mullvad_client_version: proc_macro2::TokenStream,
-    },
-}
-
-impl MullvadClient {
-    fn version(&self) -> proc_macro2::TokenStream {
-        match self {
-            MullvadClient::None {
-                mullvad_client_version,
-            } => mullvad_client_version.clone(),
-            MullvadClient::New {
-                mullvad_client_version,
-                ..
-            } => mullvad_client_version.clone(),
-        }
-    }
-}
-
-struct FunctionParameters {
-    mullvad_client: MullvadClient,
-}
-
-fn get_test_function_parameters(
-    args: &syn::punctuated::Punctuated<syn::FnArg, syn::Token![,]>,
-) -> Result<FunctionParameters> {
-    if args.len() <= 2 {
-        return Ok(FunctionParameters {
-            mullvad_client: MullvadClient::None {
-                mullvad_client_version: quote! {
-                    test_rpc::mullvad_daemon::MullvadClientVersion::None
-                },
-            },
-        });
-    }
-
-    let arg = args[2].clone();
-    let syn::FnArg::Typed(pat_type) = arg else {
-        bail!(arg, "unexpected 'mullvad_client' arg");
-    };
-
-    let syn::Type::Path(syn::TypePath { path, .. }) = &*pat_type.ty else {
-        bail!(pat_type, "unexpected 'mullvad_client' type");
-    };
-
-    let mullvad_client = match path.segments[0].ident.to_string().as_str() {
-        "mullvad_management_interface" | "MullvadProxyClient" => {
-            let mullvad_client_version =
-                quote! { test_rpc::mullvad_daemon::MullvadClientVersion::New };
-            MullvadClient::New {
-                mullvad_client_version,
-            }
-        }
-        _ => bail!(pat_type, "cannot infer mullvad client type"),
-    };
-
-    Ok(FunctionParameters { mullvad_client })
-}
diff --git a/test/test-rpc/src/mullvad_daemon.rs b/test/test-rpc/src/mullvad_daemon.rs
index e9865440b9..3a62fa7c6a 100644
--- a/test/test-rpc/src/mullvad_daemon.rs
+++ b/test/test-rpc/src/mullvad_daemon.rs
@@ -26,9 +26,3 @@ pub enum Verbosity {
     Debug,
     Trace,
 }
-
-#[derive(Clone, Copy, Debug, PartialEq)]
-pub enum MullvadClientVersion {
-    None,
-    New,
-}