fixup! WIP Add location overrides per test to test-manager config

test/test-manager/src/config/manifest.rs

@@ -55,32 +55,66 @@ impl Config {
 }
 
 mod locations {
-    use std::collections::BTreeMap;
     use std::ops::Deref;
 
     use serde::{de::Visitor, Deserialize, Serialize};
 
+    // "location": {
+    //   "override": [
+    //     { "test": "test_daita", locations: ["se-got-101"] },
+    //     { "test": "*", locations: ["Nordic"] }
+    //   ]
+    // },
+
     #[derive(Serialize, Deserialize, Clone, Default)]
     pub struct Locations {
-        pub r#override: Override,
+        pub r#override: Overrides,
     }
 
-    impl Locations {}
+    impl Locations {
+        // Look up a test (by name) and see if there are any locations
+        // that we should use.
+        pub fn lookup(&self, test: &str) -> Option<&Vec<String>> {
+            self.r#override.lookup(test)
+        }
+    }
+
+    /// Mapping of glob pattern to a set of locations.
+    #[derive(Serialize, Deserialize, Clone)]
+    pub struct Overrides(Vec<Override>);
 
     #[derive(Serialize, Deserialize, Clone)]
-    pub struct Override(BTreeMap<SerializeableGlob, Vec<String>>);
+    pub struct Override {
+        test: SerializeableGlob,
+        locations: Vec<String>,
+    }
+
+    impl Overrides {
+        // Lookup the first test that matches a glob pattern.
+        fn lookup(&self, test: &str) -> Option<&Vec<String>> {
+            self.0
+                .iter()
+                .find(
+                    |Override {
+                         test: test_glob, ..
+                     }| test_glob.matches(test),
+                )
+                .map(|Override { locations, .. }| locations)
+        }
+    }
 
-    impl Default for Override {
+    impl Default for Overrides {
         /// All tests default to using the "any" location.
         /// Written out in a config it would look like the following: { "*": ["any"] }
         fn default() -> Self {
             let overrides = {
-                let mut overrides = BTreeMap::new();
                 let glob = SerializeableGlob::from(glob::Pattern::new("*").unwrap());
-                overrides.insert(glob, vec!["any".to_string()]);
-                overrides
+                vec![Override {
+                    test: glob,
+                    locations: vec!["any".to_string()],
+                }]
             };
-            Override(overrides)
+            Overrides(overrides)
         }
     }
 

test/test-manager/src/main.rs

@@ -326,7 +326,16 @@ async fn main() -> Result<()> {
                 test_rpc::meta::Os::from(vm_config.os_type),
                 openvpn_certificate,
             ));
-            let tests = get_filtered_tests(&test_filters)?;
+            let tests = {
+                let mut tests = get_filtered_tests(&test_filters)?;
+                // Fill in location overrides
+                if let Some(locations) = &config.location {
+                    for test in tests.iter_mut() {
+                        test.location = locations.lookup(test.name).cloned();
+                    }
+                }
+                tests
+            };
 
             // For convenience, spawn a SOCKS5 server that is reachable for tests that need it
             let socks = socks_server::spawn(SocketAddr::new(

test/test-manager/src/run_tests.rs

@@ -157,13 +157,16 @@ pub async fn run(
     };
 
     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")?;
 
+        tests::prepare_custom_lists(&mut mullvad_client, &test).await?;
+
         let mullvad_client = rpc_provider
             .mullvad_client(test.mullvad_client_version)
             .await;
+
         test_handler
             .run_test(&test.func, test.name, mullvad_client)
             .await?;

test/test-manager/src/tests/helpers.rs

@@ -1344,7 +1344,7 @@ pub mod custom_lists {
 
     /// Dig out a custom list from the daemon settings based on the custom list's name.
     /// There should be an rpc for this.
-    async fn find_custom_list(
+    pub async fn find_custom_list(
         rpc: &mut MullvadProxyClient,
         name: &str,
     ) -> anyhow::Result<CustomList> {

test/test-manager/src/tests/mod.rs

@@ -16,6 +16,7 @@ mod tunnel_state;
 mod ui;
 
 use itertools::Itertools;
+use mullvad_types::relay_constraints::GeographicLocationConstraint;
 pub use test_metadata::TestMetadata;
 
 use anyhow::Context;
@@ -137,7 +138,7 @@ pub fn get_filtered_tests(specified_tests: &[String]) -> Result<Vec<TestMetadata
 pub async fn prepare_daemon(
     rpc: &ServiceClient,
     rpc_provider: &RpcClientProvider,
-) -> anyhow::Result<()> {
+) -> anyhow::Result<MullvadProxyClient> {
     // Check if daemon should be restarted
     let mut mullvad_client = ensure_daemon_version(rpc, rpc_provider)
         .await
@@ -155,6 +156,49 @@ pub async fn prepare_daemon(
     helpers::custom_lists::add_default_lists(&mut mullvad_client).await?;
     helpers::custom_lists::set_default_location(&mut mullvad_client).await?;
 
+    Ok(mullvad_client)
+}
+
+/// Create an "anonynmous" custom list for this test. The custom list will
+/// have the same as the test and contain the locations as specified by [TestMetadata::location].
+pub async fn prepare_custom_lists(
+    mullvad_client: &mut MullvadProxyClient,
+    test: &TestMetadata,
+) -> anyhow::Result<()> {
+    use helpers::custom_lists::find_custom_list;
+    // Convert locations from the test config to actual location constraints
+    let locations = {
+        let mut locations: Vec<GeographicLocationConstraint> = vec![];
+        for input in test.location.clone().unwrap() {
+            match input.parse::<GeographicLocationConstraint>() {
+                Ok(location) => locations.push(location),
+                // If some location argument can not be parsed as a GeographicLocationconstraint,
+                // assume it is a custom list.
+                Err(_parse_error) => {
+                    let custom_list = find_custom_list(mullvad_client, &input).await.unwrap();
+                    // Hack: Dig out all the geographic location constraints from this custom list
+                    for custom_list_location in custom_list.locations {
+                        locations.push(custom_list_location)
+                    }
+                }
+            };
+        }
+        locations
+    };
+
+    // Add the custom list to the current app instance
+    let id = mullvad_client
+        .create_custom_list(test.name.to_string())
+        .await?;
+
+    let mut custom_list = find_custom_list(mullvad_client, test.name).await?;
+
+    assert_eq!(id, custom_list.id);
+    for location in locations {
+        custom_list.locations.insert(location);
+    }
+    mullvad_client.update_custom_list(custom_list).await?;
+
     Ok(())
 }
 

test/test-manager/src/tests/test_metadata.rs

@@ -1,13 +1,15 @@
 use super::TestWrapperFunction;
 use test_rpc::{meta::Os, mullvad_daemon::MullvadClientVersion};
 
-#[derive(Clone)]
+#[derive(Clone, Debug)]
 pub struct TestMetadata {
     pub name: &'static str,
     pub targets: &'static [Os],
     pub mullvad_client_version: MullvadClientVersion,
     pub func: TestWrapperFunction,
     pub priority: Option<i32>,
+    // TODO: Document
+    pub location: Option<Vec<String>>,
 }
 
 // Register our test metadata struct with inventory to allow submitting tests of this type.

test/test-manager/test_macro/src/lib.rs

@@ -191,6 +191,7 @@ fn create_test(test_function: TestFunction) -> proc_macro2::TokenStream {
             mullvad_client_version: #function_mullvad_version,
             func: #wrapper_closure,
             priority: #test_function_priority,
+            location: None,
         });
     }
 }

test/test-rpc/src/mullvad_daemon.rs

@@ -27,7 +27,7 @@ pub enum Verbosity {
     Trace,
 }
 
-#[derive(Clone, Copy, PartialEq)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum MullvadClientVersion {
     None,
     New,