Merge branch 'test-add-env-flag'

6 files changed, 25 insertions, 2 deletions

diff --git a/test/connection-checker/src/cli.rs b/test/connection-checker/src/cli.rs
index 97240cf974..66402c8959 100644
--- a/test/connection-checker/src/cli.rs
+++ b/test/connection-checker/src/cli.rs
@@ -37,4 +37,7 @@ pub struct Opt {
     /// Junk data for each UDP and TCP packet
     #[clap(long, requires = "leak", default_value = "Hello there!")]
     pub payload: String,
+
+    /// URL to perform the connection check against. For example, https://am.i.mullvad.net/json.
+    pub url: String,
 }
diff --git a/test/connection-checker/src/main.rs b/test/connection-checker/src/main.rs
index c56d8b31da..6fb0d84843 100644
--- a/test/connection-checker/src/main.rs
+++ b/test/connection-checker/src/main.rs
@@ -49,7 +49,7 @@ fn am_i_mullvad(opt: &Opt) -> eyre::Result<bool> {
         mullvad_exit_ip_hostname: Option<String>,
     }
 
-    let url = "https://am.i.mullvad.net/json";
+    let url = &opt.url;
 
     let client = Client::new();
     let response: Response = client
diff --git a/test/scripts/test-utils.sh b/test/scripts/test-utils.sh
index d5f12489ea..319c646020 100755
--- a/test/scripts/test-utils.sh
+++ b/test/scripts/test-utils.sh
@@ -293,6 +293,12 @@ function run_tests_for_os {
             runner_dir_flag=()
         fi
 
+        if [ -n "${MULLVAD_HOST+x}" ]; then
+            mullvad_host_arg=("--mullvad-host" "$MULLVAD_HOST")
+        else
+            mullvad_host_arg=()
+        fi
+
         if ! RUST_LOG_STYLE=always $test_manager run-tests \
             --account "${ACCOUNT_TOKEN:?Error: ACCOUNT_TOKEN not set}" \
             --app-package "${APP_PACKAGE:?Error: APP_PACKAGE not set}" \
@@ -301,6 +307,7 @@ function run_tests_for_os {
             --package-dir "${package_dir}" \
             --vm "$vm" \
             --openvpn-certificate "${OPENVPN_CERTIFICATE:-"assets/openvpn.ca.crt"}" \
+            "${mullvad_host_arg[@]}" \
             "${test_filters_arg[@]}" \
             "${runner_dir_flag[@]}" \
             2>&1 | sed -r "s/${ACCOUNT_TOKEN}/\{ACCOUNT_TOKEN\}/g"; then
diff --git a/test/test-by-version.sh b/test/test-by-version.sh
index 337a0dfe83..e9de281bad 100755
--- a/test/test-by-version.sh
+++ b/test/test-by-version.sh
@@ -12,6 +12,7 @@ usage() {
     echo "  - APP_VERSION: The version of the app to test (defaults to the latest stable release)"
     echo "  - APP_PACKAGE_TO_UPGRADE_FROM: The package version to upgrade from (defaults to none)"
     echo "  - OPENVPN_CERTIFICATE: Path to an OpenVPN CA certificate the app should use during test (defaults to assets/openvpn.ca.crt)"
+    echo "  - MULLVAD_HOST: Conncheck and API environment to use, eg stagemole.eu (defaults to mullvad.net, or the config file if set)"
     echo "  - TEST_DIST_DIR: Relative path to a directory with prebuilt binaries as produced by scripts/build.sh."
     echo "  - TEST_FILTERS: specifies which tests to run (defaults to all)"
     echo "  - TEST_REPORT : path to save the test results in a structured format"
diff --git a/test/test-manager/src/main.rs b/test/test-manager/src/main.rs
index 88203cd8c2..8628502eff 100644
--- a/test/test-manager/src/main.rs
+++ b/test/test-manager/src/main.rs
@@ -12,7 +12,7 @@ mod vm;
 use std::{net::SocketAddr, path::PathBuf};
 
 use anyhow::{Context, Result};
-use clap::Parser;
+use clap::{builder::PossibleValuesParser, Parser};
 use tests::{config::TEST_CONFIG, get_filtered_tests};
 use vm::provision;
 
@@ -74,6 +74,11 @@ enum Commands {
         #[arg(long, group = "display_args")]
         display: bool,
 
+        /// API and conncheck environment to use. The domain name will be prefixed with "api." and
+        /// "ipv4.am.i.".
+        #[arg(long, value_parser = PossibleValuesParser::new(&["mullvad.net", "stagemole.eu", "devmole.eu"]))]
+        mullvad_host: Option<String>,
+
         /// Run VNC server on a specified port
         #[arg(long, group = "display_args")]
         vnc: Option<u16>,
@@ -236,6 +241,7 @@ async fn main() -> Result<()> {
         Commands::RunTests {
             vm,
             display,
+            mullvad_host,
             vnc,
             account,
             app_package,
@@ -256,6 +262,10 @@ async fn main() -> Result<()> {
                 (true, true) => unreachable!("invalid combination"),
             };
 
+            if let Some(mullvad_host) = mullvad_host {
+                log::trace!("Setting Mullvad host using --mullvad-host flag");
+                config.mullvad_host = Some(mullvad_host);
+            }
             let mullvad_host = config.get_host();
             log::debug!("Mullvad host: {mullvad_host}");
 
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 84945aef79..fc86a73c70 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -982,6 +982,8 @@ impl ConnChecker {
                 "--leak-tcp",
                 "--leak-udp",
                 "--leak-icmp",
+                "--url",
+                &format!("https://am.i.{}/json", TEST_CONFIG.mullvad_host),
             ]
             .map(String::from)
             .to_vec();