WIP Make daemon restart logic a bit more robust

Author: MarkusPettersson98

test/test-manager/src/run_tests.rs

@@ -102,15 +102,7 @@ pub async fn run(
         )
         .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?;
-            }
-        }
-
+        // Record test result
         if print_failed_tests_only {
             // Print results of failed test
             if matches!(test_result.result, Err(_) | Ok(Err(_))) {
@@ -137,6 +129,23 @@ pub async fn run(
         .await
         .context("Failed to log test result")?;
 
+        // Perform clean up between tests
+        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(_))) {
+                if let Err(cleanup_error) =
+                    crate::tests::cleanup_after_test(client.clone(), &test_context.rpc_provider)
+                        .await
+                {
+                    log::error!("Clean up between tests failed");
+                    log::error!("{cleanup_error}");
+                    log::error!("Exiting test run");
+                    break;
+                }
+            }
+        }
+
         match test_result.result {
             Err(panic) => {
                 failed_tests.push(test.name);

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

@@ -581,14 +581,14 @@ pub fn unreachable_wireguard_tunnel() -> talpid_types::net::wireguard::Connectio
 /// # Note
 /// This is independent of the running daemon's environment.
 /// It is solely dependant on the current value of [`TEST_CONFIG`].
-pub async fn get_app_env() -> anyhow::Result<HashMap<String, String>> {
+pub async fn get_app_env() -> Result<HashMap<String, String>, Error> {
     use mullvad_api::env;
 
     let api_host = format!("api.{}", TEST_CONFIG.mullvad_host);
     let api_host_with_port = format!("{api_host}:443");
     let api_addr = resolve_hostname_with_retries(api_host_with_port)
         .await
-        .context("failed to resolve API host")?;
+        .map_err(Error::DnsLookup)?;
 
     Ok(HashMap::from_iter(vec![
         (env::API_HOST_VAR.to_string(), api_host),

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

@@ -44,6 +44,9 @@ pub enum Error {
     #[error("geoip lookup failed")]
     GeoipLookup(#[source] test_rpc::Error),
 
+    #[error("DNS lookup failed")]
+    DnsLookup(#[source] std::io::Error),
+
     #[error("Found running daemon unexpectedly")]
     DaemonRunning,
 
@@ -76,13 +79,61 @@ pub fn get_tests() -> Vec<&'static TestMetadata> {
 
 /// Restore settings to the defaults.
 pub async fn cleanup_after_test(
-    rpc: ServiceClient,
+    mut rpc: ServiceClient,
     rpc_provider: &RpcClientProvider,
 ) -> anyhow::Result<()> {
     log::debug!("Cleaning up daemon in test cleanup");
     // Check if daemon should be restarted.
-    // TODO: The deamon needs to be up and running after this line.
-    restart_daemon(rpc).await?;
+    // TODO: Move this shizzle up one level?
+    // TODO: The daemon needs to be up and running after this line.
+    if let Err(daemon_restart_error) = restart_daemon(&rpc).await {
+        match daemon_restart_error {
+            // Something went wrong in the communication between test-manager <-> test-runner.
+            Error::Rpc(rpc_error) => {
+                log::warn!("Could not restart the daemon due to RPC-error: {rpc_error}");
+                // TODO: Try to create a new gRPC client. Need to move this logic up one level to
+                // do this.
+
+                // Restart the test-runner
+                // HACK: Accomplish this by restarting the virtual machine. This should not be
+                // necessary.
+                rpc.reboot().await.inspect_err(|_| {
+                    log::error!("Failed to reboot test runner virtual machine!");
+                })?;
+                rpc.reset_daemon_environment().await.inspect_err(|daemon_restart_failure| {
+                    log::warn!("Rebooting the test runner virtual machine did not work: {daemon_restart_failure}")
+                })?;
+            }
+            // Something wen't wrong in the daemon.
+            daemon_error @ (Error::DaemonNotRunning
+            | Error::Daemon(_)
+            | Error::UnexpectedErrorState(_)
+            | Error::ManagementInterface(_)) => {
+                log::warn!("Could not restart daemon due to daemon error: {daemon_error}");
+                log::warn!("Rebooting the test runner virtual machine");
+                // First, reboot the test-runner VM.
+                rpc.reboot().await.inspect_err(|_| {
+                    log::error!("Failed to reboot test runner virtual machine!");
+                })?;
+                if let Err(daemon_restart_failure) = rpc.reset_daemon_environment().await {
+                    log::warn!("Rebooting the test runner virtual machine did not work: {daemon_restart_failure}");
+                    log::warn!("Reinstalling the app");
+                    // TODO: If rebooting the VM did not work, try to re-install the app.
+                }
+            }
+            // We don't really care about these errors in this context.
+            non_fatal @ (Error::DaemonRunning
+            | Error::GeoipLookup(_)
+            | Error::DnsLookup(_)
+            | Error::MissingGuiTest) => {
+                log::warn!("Could not restart daemon due to non-fatal error: {non_fatal}");
+                log::warn!("Restarting dameon one more time");
+                rpc.reset_daemon_environment().await?;
+            }
+            #[cfg(target_os = "macos")]
+            Other(_) => todo!("Remove this variant, we can't handle this error properly"),
+        }
+    }
     let mut mullvad_client = rpc_provider.new_client().await;
 
     helpers::disconnect_and_wait(&mut mullvad_client).await?;
@@ -197,7 +248,7 @@ pub async fn cleanup_after_test(
 /// 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<()> {
+async fn restart_daemon(rpc: &ServiceClient) -> Result<(), Error> {
     let current_env = rpc.get_daemon_environment().await?;
     let default_env = get_app_env().await?;
     if current_env != default_env {

test/test-rpc/src/client.rs

@@ -302,6 +302,24 @@ impl ServiceClient {
         Ok(())
     }
 
+    /// Restart the Mullvad daemon with default environment variables..
+    ///
+    /// # Returns
+    /// - `Result::Ok` if the daemon was successfully restarted.
+    /// - `Result::Err(Error)` if the daemon could not be restarted and is thus no longer running.
+    pub async fn reset_daemon_environment(&self) -> Result<(), Error> {
+        let mut ctx = tarpc::context::current();
+        ctx.deadline = SystemTime::now().checked_add(LOG_LEVEL_TIMEOUT).unwrap();
+        self.client
+            .set_daemon_environment(ctx, Default::default())
+            .await??;
+
+        self.mullvad_daemon_wait_for_state(|state| state == ServiceStatus::Running)
+            .await?;
+
+        Ok(())
+    }
+
     /// Set environment variables specified by `env` and restart the Mullvad daemon.
     ///
     /// # Returns
@@ -370,6 +388,7 @@ impl ServiceClient {
         self.connection_handle.reset_connected_state().await;
         self.connection_handle.wait_for_server().await?;
 
+        // TODO: Get rid of this magic number.
         tokio::time::sleep(std::time::Duration::from_secs(5)).await;
 
         Ok(())