Remove `GetCurrentLocation`.

Make the daemon send two tunnel state updates, one with out IP being
empty, and another with it being filled when am.i.mullvad.net responds.

Update CLI for this change. Other front ends are left out.

13 files changed, 146 insertions, 161 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 02927b8545..69ee65f28d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -22,7 +22,14 @@ Line wrap the file at 100 chars.                                              Th
 * **Security**: in case of vulnerabilities.
 
 ## [Unreleased]
+### Fixed
+#### Linux
+- Out IP missing forever when am.i.mullvad.net returns error
+
 ### Changed
+- Remove `--location` flag from `mullvad status` CLI. Location and IP will now always
+  be printed (if available). `mullvad status listen` does no longer print location info.
+
 #### Android
 - Migrated to Compose Navigation
 - Allow for full rotation
diff --git a/mullvad-cli/src/cmds/status.rs b/mullvad-cli/src/cmds/status.rs
index a4a8add15e..f833868516 100644
--- a/mullvad-cli/src/cmds/status.rs
+++ b/mullvad-cli/src/cmds/status.rs
@@ -18,10 +18,7 @@ pub struct StatusArgs {
     #[arg(long, short = 'v')]
     verbose: bool,
 
-    /// Print the current location and IP, based on GeoIP lookups
-    #[arg(long, short = 'l')]
-    location: bool,
-
+    // TODO: changelog about removing location flag
     /// Enable debug output
     #[arg(long, short = 'd')]
     debug: bool,
@@ -29,22 +26,29 @@ pub struct StatusArgs {
 
 impl Status {
     pub async fn listen(mut rpc: MullvadProxyClient, args: StatusArgs) -> Result<()> {
+        let mut previous_tunnel_state = None;
+
         while let Some(event) = rpc.events_listen().await?.next().await {
             match event? {
                 DaemonEvent::TunnelState(new_state) => {
                     if args.debug {
                         println!("New tunnel state: {new_state:#?}");
                     } else {
-                        format::print_state(&new_state, args.verbose);
-                    }
-
-                    match new_state {
-                        TunnelState::Connected { .. } | TunnelState::Disconnected => {
-                            if args.location {
-                                print_location(&mut rpc).await?;
-                            }
+                        // When we enter the connected or disconnected state, am.i.mullvad.net will
+                        // be polled to get IP information. When it arrives, we will get another
+                        // tunnel state of the same enum type, but with the IP  filled in. This
+                        // match statement checks for duplicate tunnel states and skips the second
+                        // print to avoid spamming the user.
+                        match (&previous_tunnel_state, &new_state) {
+                            (Some(TunnelState::Disconnected(_)), TunnelState::Disconnected(_))
+                            | (
+                                Some(TunnelState::Connected { .. }),
+                                TunnelState::Connected { .. },
+                            ) => continue,
+                            _ => {}
                         }
-                        _ => {}
+                        format::print_state(&new_state, args.verbose);
+                        previous_tunnel_state = Some(new_state);
                     }
                 }
                 DaemonEvent::Settings(settings) => {
@@ -88,11 +92,9 @@ pub async fn handle(cmd: Option<Status>, args: StatusArgs) -> Result<()> {
     if args.debug {
         println!("Tunnel state: {state:#?}");
     } else {
+        // TODO: respect location arg?
         format::print_state(&state, args.verbose);
-    }
-
-    if args.location {
-        print_location(&mut rpc).await?;
+        format::print_location(&state);
     }
 
     if cmd == Some(Status::Listen) {
@@ -101,31 +103,6 @@ pub async fn handle(cmd: Option<Status>, args: StatusArgs) -> Result<()> {
     Ok(())
 }
 
-async fn print_location(rpc: &mut MullvadProxyClient) -> Result<()> {
-    let location = match rpc.get_current_location().await {
-        Ok(location) => location,
-        Err(error) => match &error {
-            mullvad_management_interface::Error::NoLocationData => {
-                println!("Location data unavailable");
-                return Ok(());
-            }
-            _ => return Err(error.into()),
-        },
-    };
-    if let Some(ipv4) = location.ipv4 {
-        println!("IPv4: {ipv4}");
-    }
-    if let Some(ipv6) = location.ipv6 {
-        println!("IPv6: {ipv6}");
-    }
-
-    println!(
-        "Position: {:.5}°N, {:.5}°W",
-        location.latitude, location.longitude
-    );
-    Ok(())
-}
-
 fn print_account_loggedout(state: &TunnelState, device: &DeviceState) {
     match state {
         TunnelState::Connecting { .. } | TunnelState::Connected { .. } | TunnelState::Error(_) => {
@@ -137,6 +114,6 @@ fn print_account_loggedout(state: &TunnelState, device: &DeviceState) {
                 DeviceState::LoggedIn(_) => (),
             }
         }
-        TunnelState::Disconnected | TunnelState::Disconnecting(_) => (),
+        TunnelState::Disconnected(_) | TunnelState::Disconnecting(_) => (),
     }
 }
diff --git a/mullvad-cli/src/format.rs b/mullvad-cli/src/format.rs
index 2938842cef..0e920b6218 100644
--- a/mullvad-cli/src/format.rs
+++ b/mullvad-cli/src/format.rs
@@ -37,11 +37,34 @@ pub fn print_state(state: &TunnelState, verbose: bool) {
                 format_relay_connection(endpoint, location.as_ref(), verbose)
             );
         }
-        Disconnected => println!("Disconnected"),
+        Disconnected(_) => {
+            println!("Disconnected");
+        }
         Disconnecting(_) => println!("Disconnecting..."),
     }
 }
 
+pub fn print_location(state: &TunnelState) {
+    let location = match state {
+        TunnelState::Disconnected(location) => location,
+        TunnelState::Connected { location, .. } => location,
+        _ => return,
+    };
+    if let Some(location) = location {
+        print!("Your connection appears from: {}", location.country);
+        if let Some(city) = &location.city {
+            print!(", {}", city);
+        }
+        if let Some(ipv4) = location.ipv4 {
+            print!(". IPv4: {ipv4}");
+        }
+        if let Some(ipv6) = location.ipv6 {
+            print!(", IPv6: {ipv6}");
+        }
+        println!();
+    }
+}
+
 fn format_relay_connection(
     endpoint: &TunnelEndpoint,
     location: Option<&GeoIpLocation>,
diff --git a/mullvad-daemon/src/geoip.rs b/mullvad-daemon/src/geoip.rs
index 0af495e8bd..b4ed662fa1 100644
--- a/mullvad-daemon/src/geoip.rs
+++ b/mullvad-daemon/src/geoip.rs
@@ -44,13 +44,13 @@ const LOCATION_RETRY_STRATEGY: Jittered<ExponentialBackoff> =
     Jittered::jitter(ExponentialBackoff::new(Duration::from_secs(1), 4));
 
 /// Fetch the current `GeoIpLocation` with retrys
-pub async fn get_geo_location(
+pub async fn get_geo_location_with_retry(
     rest_service: RequestServiceHandle,
     use_ipv6: bool,
     api_handle: ApiAvailabilityHandle,
-) -> Option<GeoIpLocation> {
+) -> Result<GeoIpLocation, Error> {
     log::debug!("Fetching GeoIpLocation");
-    match retry_future(
+    retry_future(
         move || send_location_request(rest_service.clone(), use_ipv6),
         move |result| match result {
             Err(error) if error.is_network_error() => !api_handle.get_state().is_offline(),
@@ -59,13 +59,6 @@ pub async fn get_geo_location(
         LOCATION_RETRY_STRATEGY,
     )
     .await
-    {
-        Ok(loc) => Some(loc),
-        Err(e) => {
-            log::warn!("Unable to fetch GeoIP location: {}", e.display_chain());
-            None
-        }
-    }
 }
 
 async fn send_location_request(
diff --git a/mullvad-daemon/src/lib.rs b/mullvad-daemon/src/lib.rs
index 58069db4fa..6c78b3452d 100644
--- a/mullvad-daemon/src/lib.rs
+++ b/mullvad-daemon/src/lib.rs
@@ -27,7 +27,7 @@ mod tunnel;
 pub mod version;
 mod version_check;
 
-use crate::{geoip::get_geo_location, target_state::PersistentTargetState};
+use crate::target_state::PersistentTargetState;
 use device::{AccountEvent, PrivateAccountAndDevice, PrivateDeviceEvent};
 use futures::{
     channel::{mpsc, oneshot},
@@ -203,8 +203,6 @@ pub enum DaemonCommand {
     Reconnect(oneshot::Sender<bool>),
     /// Request the current state.
     GetState(oneshot::Sender<TunnelState>),
-    /// Get the current geographical location.
-    GetCurrentLocation(oneshot::Sender<Option<GeoIpLocation>>),
     CreateNewAccount(ResponseTx<String, Error>),
     /// Request the metadata for an account.
     GetAccountData(
@@ -420,7 +418,7 @@ impl DaemonExecutionState {
         match self {
             Running => {
                 match tunnel_state {
-                    TunnelState::Disconnected => mem::replace(self, Finished),
+                    TunnelState::Disconnected(_) => mem::replace(self, Finished),
                     _ => mem::replace(self, Exiting),
                 };
             }
@@ -617,6 +615,7 @@ pub struct Daemon<L: EventListener> {
     tunnel_state_machine_handle: TunnelStateMachineHandle,
     #[cfg(target_os = "windows")]
     volume_update_tx: mpsc::UnboundedSender<()>,
+    location_abort_handle: Option<tokio::task::AbortHandle>,
 }
 
 impl<L> Daemon<L>
@@ -831,7 +830,7 @@ where
         relay_list_updater.update().await;
 
         let daemon = Daemon {
-            tunnel_state: TunnelState::Disconnected,
+            tunnel_state: TunnelState::Disconnected(None),
             target_state,
             state: DaemonExecutionState::Running,
             #[cfg(target_os = "linux")]
@@ -857,6 +856,7 @@ where
             tunnel_state_machine_handle,
             #[cfg(target_os = "windows")]
             volume_update_tx,
+            location_abort_handle: None,
         };
 
         api_availability.unsuspend();
@@ -960,15 +960,26 @@ where
             .handle_state_transition(&tunnel_state_transition);
 
         let tunnel_state = match tunnel_state_transition {
-            TunnelStateTransition::Disconnected => TunnelState::Disconnected,
+            TunnelStateTransition::Disconnected => {
+                self.notify_tunnel_state_when_ip_arrives(None, TunnelState::Disconnected)
+                    .await;
+                TunnelState::Disconnected(None)
+            }
             TunnelStateTransition::Connecting(endpoint) => TunnelState::Connecting {
                 endpoint,
                 location: self.parameters_generator.get_last_location().await,
             },
-            TunnelStateTransition::Connected(endpoint) => TunnelState::Connected {
-                endpoint,
-                location: self.parameters_generator.get_last_location().await,
-            },
+            TunnelStateTransition::Connected(endpoint) => {
+                let location = self.parameters_generator.get_last_location().await;
+                let make_tunnel_state = |location| TunnelState::Connected { endpoint, location };
+                self.notify_tunnel_state_when_ip_arrives(
+                    location.clone(),
+                    make_tunnel_state.clone(),
+                )
+                .await;
+
+                make_tunnel_state(location)
+            }
             TunnelStateTransition::Disconnecting(after_disconnect) => {
                 TunnelState::Disconnecting(after_disconnect)
             }
@@ -985,7 +996,7 @@ where
         log::debug!("New tunnel state: {:?}", tunnel_state);
 
         match tunnel_state {
-            TunnelState::Disconnected => {
+            TunnelState::Disconnected(_) => {
                 self.api_handle.availability.reset_inactivity_timer();
             }
             _ => {
@@ -994,7 +1005,7 @@ where
         }
 
         match &tunnel_state {
-            TunnelState::Disconnected => self.state.disconnected(),
+            TunnelState::Disconnected(_) => self.state.disconnected(),
             TunnelState::Connecting { .. } => {
                 log::debug!("Settings: {}", self.settings.summary());
             }
@@ -1024,6 +1035,53 @@ where
         self.event_listener.notify_new_state(tunnel_state);
     }
 
+    /// Get the out IP from am.i.mullvad.net. When it arrives, send another
+    /// TunnelState with the IP filled in.
+    async fn notify_tunnel_state_when_ip_arrives(
+        &mut self,
+        old_location: Option<GeoIpLocation>,
+        make_tunnel_state: impl FnOnce(Option<GeoIpLocation>) -> TunnelState + Send + 'static,
+    ) {
+        // Abort any ongoing calls to `self.notify_tunnel_state_when_ip_arrives` from
+        // previous tunnel state transitions to avoid sending an outdated information.
+        if let Some(handle) = self.location_abort_handle.take() {
+            handle.abort();
+        }
+        let rest_service = self.api_runtime.rest_handle().await;
+        let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
+        let api_handle = self.api_handle.availability.clone();
+        let even_listener = self.event_listener.clone();
+
+        self.location_abort_handle = Some(
+            tokio::spawn(async move {
+                let merged_location =
+                    geoip::get_geo_location_with_retry(rest_service, use_ipv6, api_handle)
+                        .await
+                        .ok()
+                        // Replace the old location with new information from am.i.mullvad.net,
+                        // but keep the hostname as it is always none
+                        .map(|new_location| match old_location {
+                            Some(old_location) => GeoIpLocation {
+                                hostname: old_location.hostname,
+                                bridge_hostname: old_location.bridge_hostname,
+                                entry_hostname: old_location.entry_hostname,
+                                obfuscator_hostname: old_location.obfuscator_hostname,
+                                ..new_location
+                            },
+                            None => GeoIpLocation {
+                                hostname: None,
+                                bridge_hostname: None,
+                                entry_hostname: None,
+                                obfuscator_hostname: None,
+                                ..new_location
+                            },
+                        });
+                even_listener.notify_new_state(make_tunnel_state(merged_location));
+            })
+            .abort_handle(),
+        );
+    }
+
     fn reset_rpc_sockets_on_tunnel_state_transition(
         &mut self,
         tunnel_state_transition: &TunnelStateTransition,
@@ -1075,7 +1133,6 @@ where
             SetTargetState(tx, state) => self.on_set_target_state(tx, state).await,
             Reconnect(tx) => self.on_reconnect(tx),
             GetState(tx) => self.on_get_state(tx),
-            GetCurrentLocation(tx) => self.on_get_current_location(tx).await,
             CreateNewAccount(tx) => self.on_create_new_account(tx),
             GetAccountData(tx, account_token) => self.on_get_account_data(tx, account_token),
             GetWwwAuthToken(tx) => self.on_get_www_auth_token(tx).await,
@@ -1326,52 +1383,6 @@ where
         Self::oneshot_send(tx, performing_post_upgrade, "performing post upgrade");
     }
 
-    async fn on_get_current_location(&mut self, tx: oneshot::Sender<Option<GeoIpLocation>>) {
-        use self::TunnelState::*;
-
-        match &self.tunnel_state {
-            Disconnected => self.update_and_send_geo_location(tx, None).await,
-            Connecting { location, .. } => {
-                Self::oneshot_send(tx, location.clone(), "current location")
-            }
-            Disconnecting(..) => Self::oneshot_send(
-                tx,
-                self.parameters_generator.get_last_location().await,
-                "current location",
-            ),
-            Connected { location, .. } => {
-                self.update_and_send_geo_location(tx, location.clone())
-                    .await
-            }
-            // We are not online at all at this stage so no location data is available.
-            Error(_) => Self::oneshot_send(tx, None, "current location"),
-        };
-    }
-
-    /// Fetch the current `GeoIpLocation` and send it on the return channel,
-    /// in a non-blocking fashion. Optionally give a chached previous location.
-    async fn update_and_send_geo_location(
-        &mut self,
-        tx: oneshot::Sender<Option<GeoIpLocation>>,
-        current_relay_location: Option<GeoIpLocation>,
-    ) {
-        let rest_service = self.api_runtime.rest_handle().await;
-        let use_ipv6 = self.settings.tunnel_options.generic.enable_ipv6;
-        let api_handle = self.api_handle.availability.clone();
-        tokio::spawn(async move {
-            let new_location = get_geo_location(rest_service, use_ipv6, api_handle).await;
-            Self::oneshot_send(
-                tx,
-                new_location.map(|fetched_location| GeoIpLocation {
-                    ipv4: fetched_location.ipv4,
-                    ipv6: fetched_location.ipv6,
-                    ..current_relay_location.unwrap_or(fetched_location)
-                }),
-                "current location",
-            );
-        });
-    }
-
     fn on_create_new_account(&mut self, tx: ResponseTx<String, Error>) {
         let account_manager = self.account_manager.clone();
         tokio::spawn(async move {
diff --git a/mullvad-daemon/src/management_interface.rs b/mullvad-daemon/src/management_interface.rs
index c194825a34..cf007e665d 100644
--- a/mullvad-daemon/src/management_interface.rs
+++ b/mullvad-daemon/src/management_interface.rs
@@ -191,17 +191,6 @@ impl ManagementService for ManagementServiceImpl {
             .map(|relays| Response::new(types::RelayList::from(relays)))
     }
 
-    async fn get_current_location(&self, _: Request<()>) -> ServiceResult<types::GeoIpLocation> {
-        log::debug!("get_current_location");
-        let (tx, rx) = oneshot::channel();
-        self.send_command_to_daemon(DaemonCommand::GetCurrentLocation(tx))?;
-        let result = self.wait_for_result(rx).await?;
-        match result {
-            Some(geoip) => Ok(Response::new(types::GeoIpLocation::from(geoip))),
-            None => Err(Status::not_found("no location was found")),
-        }
-    }
-
     async fn set_bridge_settings(
         &self,
         request: Request<types::BridgeSettings>,
diff --git a/mullvad-management-interface/proto/management_interface.proto b/mullvad-management-interface/proto/management_interface.proto
index d66707b79f..97507ca8a2 100644
--- a/mullvad-management-interface/proto/management_interface.proto
+++ b/mullvad-management-interface/proto/management_interface.proto
@@ -28,7 +28,6 @@ service ManagementService {
   // Relays and tunnel constraints
   rpc UpdateRelayLocations(google.protobuf.Empty) returns (google.protobuf.Empty) {}
   rpc GetRelayLocations(google.protobuf.Empty) returns (RelayList) {}
-  rpc GetCurrentLocation(google.protobuf.Empty) returns (GeoIpLocation) {}
   rpc SetRelaySettings(RelaySettings) returns (google.protobuf.Empty) {}
   rpc SetBridgeSettings(BridgeSettings) returns (google.protobuf.Empty) {}
   rpc SetBridgeState(BridgeState) returns (google.protobuf.Empty) {}
@@ -172,7 +171,7 @@ message ErrorState {
 }
 
 message TunnelState {
-  message Disconnected {}
+  message Disconnected { GeoIpLocation disconnected_location = 1; }
   message Connecting { TunnelStateRelayInfo relay_info = 1; }
   message Connected { TunnelStateRelayInfo relay_info = 1; }
   message Disconnecting { AfterDisconnect after_disconnect = 1; }
diff --git a/mullvad-management-interface/src/client.rs b/mullvad-management-interface/src/client.rs
index 1c9d80b2e8..85048e29e7 100644
--- a/mullvad-management-interface/src/client.rs
+++ b/mullvad-management-interface/src/client.rs
@@ -7,7 +7,6 @@ use mullvad_types::{
     account::{AccountData, AccountToken, VoucherSubmission},
     custom_list::{CustomList, Id},
     device::{Device, DeviceEvent, DeviceId, DeviceState, RemoveDeviceEvent},
-    location::GeoIpLocation,
     relay_constraints::{
         BridgeSettings, BridgeState, ObfuscationSettings, RelayOverride, RelaySettings,
     },
@@ -234,16 +233,6 @@ impl MullvadProxyClient {
         Ok(())
     }
 
-    pub async fn get_current_location(&mut self) -> Result<GeoIpLocation> {
-        let location = self
-            .0
-            .get_current_location(())
-            .await
-            .map_err(map_location_error)?
-            .into_inner();
-        GeoIpLocation::try_from(location).map_err(Error::InvalidResponse)
-    }
-
     pub async fn set_bridge_settings(&mut self, settings: BridgeSettings) -> Result<()> {
         let settings = types::BridgeSettings::from(settings);
         self.0
@@ -571,10 +560,9 @@ impl MullvadProxyClient {
     /// Set the [`AccessMethod`] which [`ApiConnectionModeProvider`] should
     /// pick.
     ///
-    /// - `access_method`: If `Some(access_method)`, [`ApiConnectionModeProvider`] will skip
-    ///     ahead and return `access_method` when asked for a new access method.
-    ///     If `None`, [`ApiConnectionModeProvider`] will pick the next access
-    ///     method "randomly"
+    /// - `access_method`: If `Some(access_method)`, [`ApiConnectionModeProvider`] will skip ahead
+    ///   and return `access_method` when asked for a new access method. If `None`,
+    ///   [`ApiConnectionModeProvider`] will pick the next access method "randomly"
     ///
     /// [`ApiConnectionModeProvider`]: mullvad_daemon::api::ApiConnectionModeProvider
     pub async fn set_access_method(&mut self, api_access_method: access_method::Id) -> Result<()> {
@@ -696,13 +684,6 @@ fn map_device_error(status: Status) -> Error {
     }
 }
 
-fn map_location_error(status: Status) -> Error {
-    match status.code() {
-        Code::NotFound => Error::NoLocationData,
-        _other => Error::Rpc(status),
-    }
-}
-
 fn map_custom_list_error(status: Status) -> Error {
     match status.code() {
         Code::NotFound => {
diff --git a/mullvad-management-interface/src/types/conversions/states.rs b/mullvad-management-interface/src/types/conversions/states.rs
index b742aa7b00..2541859e27 100644
--- a/mullvad-management-interface/src/types/conversions/states.rs
+++ b/mullvad-management-interface/src/types/conversions/states.rs
@@ -32,8 +32,10 @@ impl From<mullvad_types::states::TunnelState> for proto::TunnelState {
             };
 
         let state = match state {
-            MullvadTunnelState::Disconnected => {
-                proto::tunnel_state::State::Disconnected(proto::tunnel_state::Disconnected {})
+            MullvadTunnelState::Disconnected(disconnected_location) => {
+                proto::tunnel_state::State::Disconnected(proto::tunnel_state::Disconnected {
+                    disconnected_location: disconnected_location.map(proto::GeoIpLocation::from),
+                })
             }
             MullvadTunnelState::Connecting { endpoint, location } => {
                 proto::tunnel_state::State::Connecting(proto::tunnel_state::Connecting {
@@ -189,7 +191,13 @@ impl TryFrom<proto::TunnelState> for mullvad_types::states::TunnelState {
         use talpid_types::{net as talpid_net, tunnel as talpid_tunnel};
 
         let state = match state.state {
-            Some(proto::tunnel_state::State::Disconnected(_)) => MullvadState::Disconnected,
+            Some(proto::tunnel_state::State::Disconnected(proto::tunnel_state::Disconnected {
+                disconnected_location,
+            })) => MullvadState::Disconnected(
+                disconnected_location
+                    .map(mullvad_types::location::GeoIpLocation::try_from)
+                    .transpose()?,
+            ),
             Some(proto::tunnel_state::State::Connecting(proto::tunnel_state::Connecting {
                 relay_info:
                     Some(proto::TunnelStateRelayInfo {
diff --git a/mullvad-types/src/states.rs b/mullvad-types/src/states.rs
index f6131711cb..4a266d1d6e 100644
--- a/mullvad-types/src/states.rs
+++ b/mullvad-types/src/states.rs
@@ -34,7 +34,7 @@ impl fmt::Display for TargetState {
 #[cfg_attr(target_os = "android", derive(IntoJava))]
 #[cfg_attr(target_os = "android", jnix(package = "net.mullvad.mullvadvpn.model"))]
 pub enum TunnelState {
-    Disconnected,
+    Disconnected(Option<GeoIpLocation>),
     Connecting {
         endpoint: TunnelEndpoint,
         location: Option<GeoIpLocation>,
@@ -60,6 +60,6 @@ impl TunnelState {
 
     /// Returns true if the tunnel state is in the disconnected state.
     pub fn is_disconnected(&self) -> bool {
-        matches!(self, TunnelState::Disconnected)
+        matches!(self, TunnelState::Disconnected(_))
     }
 }
diff --git a/test/test-manager/src/tests/helpers.rs b/test/test-manager/src/tests/helpers.rs
index 48dba20aac..3f64aef1be 100644
--- a/test/test-manager/src/tests/helpers.rs
+++ b/test/test-manager/src/tests/helpers.rs
@@ -229,7 +229,7 @@ pub async fn disconnect_and_wait(
         .await
         .map_err(|error| Error::DaemonError(format!("failed to begin disconnecting: {}", error)))?;
     wait_for_tunnel_state(mullvad_client.clone(), |state| {
-        matches!(state, TunnelState::Disconnected)
+        matches!(state, TunnelState::Disconnected(_))
     })
     .await?;
 
diff --git a/test/test-manager/src/tests/settings.rs b/test/test-manager/src/tests/settings.rs
index aea81028fe..d2bde57fcc 100644
--- a/test/test-manager/src/tests/settings.rs
+++ b/test/test-manager/src/tests/settings.rs
@@ -107,9 +107,8 @@ pub async fn test_lockdown(
     let inet_destination: SocketAddr = "1.1.1.1:1337".parse().unwrap();
 
     log::info!("Verify tunnel state: disconnected");
-    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected);
+    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected(_));
 
-    //
     // Enable lockdown mode
     //
     mullvad_client
diff --git a/test/test-manager/src/tests/tunnel_state.rs b/test/test-manager/src/tests/tunnel_state.rs
index 9ae817c7ee..d8cce41b3a 100644
--- a/test/test-manager/src/tests/tunnel_state.rs
+++ b/test/test-manager/src/tests/tunnel_state.rs
@@ -32,9 +32,8 @@ pub async fn test_disconnected_state(
     let inet_destination = "1.3.3.7:1337".parse().unwrap();
 
     log::info!("Verify tunnel state: disconnected");
-    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected);
+    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected(_));
 
-    //
     // Test whether outgoing packets can be observed
     //
 
@@ -90,7 +89,7 @@ pub async fn test_connecting_state(
     let lan_dns: SocketAddr = SocketAddr::new(IpAddr::V4(DUMMY_LAN_INTERFACE_IP), 53);
 
     log::info!("Verify tunnel state: disconnected");
-    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected);
+    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected(_));
 
     let relay_settings = RelaySettings::CustomTunnelEndpoint(CustomTunnelEndpoint {
         host: "1.3.3.7".to_owned(),
@@ -172,9 +171,8 @@ pub async fn test_error_state(
     let lan_dns: SocketAddr = SocketAddr::new(IpAddr::V4(DUMMY_LAN_INTERFACE_IP), 53);
 
     log::info!("Verify tunnel state: disconnected");
-    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected);
+    assert_tunnel_state!(&mut mullvad_client, TunnelState::Disconnected(_));
 
-    //
     // Connect to non-existent location
     //