Add TunnelMonitorStateTests

5 files changed, 159 insertions, 3 deletions

diff --git a/ios/MullvadVPN.xcodeproj/project.pbxproj b/ios/MullvadVPN.xcodeproj/project.pbxproj
index 4bdf489720..0d8656bbd1 100644
--- a/ios/MullvadVPN.xcodeproj/project.pbxproj
+++ b/ios/MullvadVPN.xcodeproj/project.pbxproj
@@ -744,6 +744,7 @@
 		A9173C372C36CD2B00F6A08C /* MullvadTypes.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 58D223D5294C8E5E0029F5F8 /* MullvadTypes.framework */; platformFilter = ios; };
 		A91D78E42B03C01600FCD5D3 /* MullvadSettings.framework in Frameworks */ = {isa = PBXBuildFile; fileRef = 58B2FDD32AA71D2A003EB5C6 /* MullvadSettings.framework */; };
 		A91EBEDA2C1337040004A84D /* RetryStrategyTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A91EBED92C1337040004A84D /* RetryStrategyTests.swift */; };
+		A923D1212D5A40A80066C090 /* TunnelMonitorStateTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A923D1202D5A40A80066C090 /* TunnelMonitorStateTests.swift */; };
 		A93181A12B727ED700E341D2 /* TunnelSettingsV4.swift in Sources */ = {isa = PBXBuildFile; fileRef = A93181A02B727ED700E341D2 /* TunnelSettingsV4.swift */; };
 		A932D9EF2B5ADD0700999395 /* ProxyConfigurationTransportProvider.swift in Sources */ = {isa = PBXBuildFile; fileRef = A932D9EE2B5ADD0700999395 /* ProxyConfigurationTransportProvider.swift */; };
 		A932D9F32B5EB61100999395 /* HeadRequestTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = A932D9F22B5EB61100999395 /* HeadRequestTests.swift */; };
@@ -2148,6 +2149,7 @@
 		A91614D02B108D1B00F416EB /* TransportLayer.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TransportLayer.swift; sourceTree = "<group>"; };
 		A917352029FAAA5200D5DCFD /* TransportStrategyTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TransportStrategyTests.swift; sourceTree = "<group>"; };
 		A91EBED92C1337040004A84D /* RetryStrategyTests.swift */ = {isa = PBXFileReference; fileEncoding = 4; lastKnownFileType = sourcecode.swift; path = RetryStrategyTests.swift; sourceTree = "<group>"; };
+		A923D1202D5A40A80066C090 /* TunnelMonitorStateTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelMonitorStateTests.swift; sourceTree = "<group>"; };
 		A92962582B1F4FDB00DFB93B /* PrivacyInfo.xcprivacy */ = {isa = PBXFileReference; lastKnownFileType = text.xml; path = PrivacyInfo.xcprivacy; sourceTree = "<group>"; };
 		A92ECC202A77FFAF0052F1B1 /* TunnelSettings.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TunnelSettings.swift; sourceTree = "<group>"; };
 		A92ECC232A7802520052F1B1 /* StoredAccountData.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = StoredAccountData.swift; sourceTree = "<group>"; };
@@ -3620,6 +3622,7 @@
 				5838321C2AC1C54600EA2071 /* TaskSleepTests.swift */,
 				58092E532A8B832E00C3CC72 /* TunnelMonitorTests.swift */,
 				F062B94C2C16E09700B6D47A /* TunnelSettingsManagerTests.swift */,
+				A923D1202D5A40A80066C090 /* TunnelMonitorStateTests.swift */,
 			);
 			path = PacketTunnelCoreTests;
 			sourceTree = "<group>";
@@ -5899,6 +5902,7 @@
 				58FE25EE2AA7764E003D1918 /* TunnelAdapterDummy.swift in Sources */,
 				581F23AD2A8CF92100788AB6 /* DefaultPathObserverFake.swift in Sources */,
 				F07751582C50F149006E6A12 /* MultiHopEphemeralPeerExchangerTests.swift in Sources */,
+				A923D1212D5A40A80066C090 /* TunnelMonitorStateTests.swift in Sources */,
 				5838321B2AC1B18400EA2071 /* PacketTunnelActor+Mocks.swift in Sources */,
 				5838321D2AC1C54600EA2071 /* TaskSleepTests.swift in Sources */,
 				58092E542A8B832E00C3CC72 /* TunnelMonitorTests.swift in Sources */,
diff --git a/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitor.swift b/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitor.swift
index 8256bd6f16..6d28c49c1d 100644
--- a/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitor.swift
+++ b/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitor.swift
@@ -269,6 +269,7 @@ public actor TunnelMonitorActor: TunnelMonitorProtocol {
     }
 
     #if DEBUG
+    /// Helper function used to help the state pass across the actor's isolation region
     internal func getState() -> TunnelMonitorState {
         TunnelMonitorState(
             connectionState: state.connectionState,
diff --git a/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitorState.swift b/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitorState.swift
index 0ec255e7fa..252a50789a 100644
--- a/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitorState.swift
+++ b/ios/PacketTunnelCore/TunnelMonitor/TunnelMonitorState.swift
@@ -178,7 +178,7 @@ public struct TunnelMonitorState {
 
         let timeSinceLastPing = now.timeIntervalSince(lastRequestDate)
         if let lastReplyDate = pingStats.lastReplyDate,
-           lastRequestDate.timeIntervalSince(lastReplyDate) >= timings.heartbeatReplyTimeout,
+           lastReplyDate.timeIntervalSince(lastRequestDate) >= timings.heartbeatReplyTimeout,
            timeSinceLastPing >= timings.pingDelay, !isHeartbeatSuspended {
             return .retryHeartbeatPing
         }
diff --git a/ios/PacketTunnelCoreTests/TunnelMonitorStateTests.swift b/ios/PacketTunnelCoreTests/TunnelMonitorStateTests.swift
new file mode 100644
index 0000000000..27d1847b48
--- /dev/null
+++ b/ios/PacketTunnelCoreTests/TunnelMonitorStateTests.swift
@@ -0,0 +1,151 @@
+//
+//  TunnelMonitorStateTests.swift
+//  PacketTunnelCoreTests
+//
+//  Created by Marco Nikic on 2025-02-10.
+//  Copyright © 2025 Mullvad VPN AB. All rights reserved.
+//
+
+import MullvadTypes
+@testable import PacketTunnelCore
+import Testing
+
+struct TunnelMonitorStateTests {
+    var testDateComponents = DateComponents()
+    var defaultPingTimeout: Duration = .milliseconds(500)
+    var hearbeatReplyTimeout: Duration = .seconds(1)
+    var pingDelay: Duration = .seconds(1)
+
+    init() async throws {
+        testDateComponents.day = 10
+        testDateComponents.month = 2
+        testDateComponents.year = 2025
+        testDateComponents.hour = 10
+        testDateComponents.calendar = .current
+    }
+
+    @Test(arguments: [
+        TunnelMonitorConnectionState.stopped,
+        TunnelMonitorConnectionState.pendingStart,
+        TunnelMonitorConnectionState.recovering,
+        TunnelMonitorConnectionState
+            .waitingConnectivity,
+    ])
+    func connectionIsOk(initialState: TunnelMonitorConnectionState) async throws {
+        let state = createState(initialState)
+        #expect(state.evaluateConnection(now: .now, pingTimeout: .zero) == .ok)
+    }
+
+    @Test(arguments: [TunnelMonitorConnectionState.connecting, TunnelMonitorConnectionState.connected])
+    func timeoutWhenReplyComesAfterPingTimeoutIn(state: TunnelMonitorConnectionState) async throws {
+        var state = createState(state)
+
+        let now = try #require(testDateComponents.date)
+        let oneSecondAgo = try #require(testDateComponents.date?.addingTimeInterval(-1))
+
+        // Sent ping a second ago
+        state.updatePingStats(sendResult: PingerSendResult(sequenceNumber: 1), now: oneSecondAgo)
+        // 0 latency reply
+        let timestamp = state.setPingReplyReceived(1, now: oneSecondAgo)
+        #expect(timestamp == oneSecondAgo)
+        #expect(state.evaluateConnection(now: now, pingTimeout: state.timings.pingTimeout) == .pingTimeout)
+    }
+
+    @Test(arguments: [TunnelMonitorConnectionState.connecting, TunnelMonitorConnectionState.connected])
+    func evaluateStateBeforeInitialPingIsSent(state: TunnelMonitorConnectionState) async throws {
+        let state = createState(state)
+        let now = try #require(testDateComponents.date)
+        #expect(state.evaluateConnection(now: now, pingTimeout: state.timings.pingTimeout) == .sendInitialPing)
+    }
+
+    @Test func evaluateConnectingStateAfterLastRequestWithoutReply() async throws {
+        var state = createState(.connecting)
+
+        let now = try #require(testDateComponents.date)
+        let nextPingDelay = state.timings.pingDelay + .seconds(1)
+        let oneSecondAgo = try #require(testDateComponents.date?.addingTimeInterval(-nextPingDelay.timeInterval))
+
+        // Sent ping a second ago
+        state.updatePingStats(sendResult: PingerSendResult(sequenceNumber: 1), now: oneSecondAgo)
+
+        #expect(state.evaluateConnection(now: now, pingTimeout: .seconds(500)) == .sendNextPing)
+    }
+
+    @Test func retryHeartbeatPingWhenHeartbeatNotSuspended() async throws {
+        var state = createState(.connected)
+
+        // Send a ping and acknowledge it
+        let oneSecondAgo = try #require(testDateComponents.date?.addingTimeInterval(-1))
+        let now = try #require(testDateComponents.date)
+        state.updatePingStats(sendResult: PingerSendResult(sequenceNumber: 1), now: oneSecondAgo)
+        _ = state.setPingReplyReceived(1, now: now)
+
+        #expect(state.evaluateConnection(now: now, pingTimeout: .hours(1)) == .retryHeartbeatPing)
+    }
+
+    @Test mutating func witnessTrafficFlowingAsExpected() async throws {
+        // Ignore heartbeat
+        hearbeatReplyTimeout = .hours(1)
+        var state = createState(.connected)
+
+        // Send a ping and acknowledge it
+        let oneSecondAgo = try #require(testDateComponents.date?.addingTimeInterval(-1))
+        let now = try #require(testDateComponents.date)
+        state.updatePingStats(sendResult: PingerSendResult(sequenceNumber: 1), now: oneSecondAgo)
+        _ = state.setPingReplyReceived(1, now: now)
+
+        #expect(state.evaluateConnection(now: now, pingTimeout: .hours(1)) == .ok)
+    }
+
+    @Test mutating func sendHeartbeatPingWhenConnected() async throws {
+        hearbeatReplyTimeout = .milliseconds(500)
+        var state = createState(.connected)
+        state.isHeartbeatSuspended = true
+
+        // Send a ping and acknowledge it
+        let tenSecondAgo = try #require(testDateComponents.date?.addingTimeInterval(-10))
+        let now = try #require(testDateComponents.date)
+        state.updatePingStats(sendResult: PingerSendResult(sequenceNumber: 1), now: tenSecondAgo)
+        _ = state.setPingReplyReceived(1, now: now)
+        let newNetStats = WgStats(bytesReceived: 100, bytesSent: 100)
+        state.updateNetStats(newStats: newNetStats, now: now)
+
+        #expect(state.evaluateConnection(now: now, pingTimeout: .hours(1)) == .sendHeartbeatPing)
+    }
+
+    @Test(arguments: [
+        TunnelMonitorConnectionState.stopped,
+        TunnelMonitorConnectionState.connected,
+        TunnelMonitorConnectionState.pendingStart,
+        TunnelMonitorConnectionState.recovering,
+        TunnelMonitorConnectionState
+            .waitingConnectivity,
+    ])
+
+    func pingTimeoutIsUnalteredIn(state: TunnelMonitorConnectionState) async throws {
+        let state = createState(state)
+        #expect(state.getPingTimeout() == defaultPingTimeout)
+    }
+
+    func createState(_ initialState: TunnelMonitorConnectionState) -> TunnelMonitorState {
+        let timings = TunnelMonitorTimings(
+            heartbeatReplyTimeout: hearbeatReplyTimeout,
+            pingTimeout: defaultPingTimeout,
+            pingDelay: pingDelay,
+            initialEstablishTimeout: .milliseconds(50),
+            connectivityCheckInterval: .milliseconds(10)
+        )
+
+        return TunnelMonitorState(
+            connectionState: initialState,
+            netStats: WgStats(),
+            pingStats: PingStats(),
+            timeoutReference: Date(),
+            lastSeenRx: nil,
+            lastSeenTx: nil,
+            isHeartbeatSuspended: false,
+            retryAttempt: 0,
+            timings: timings
+        )
+    }
+}
diff --git a/ios/PacketTunnelCoreTests/TunnelMonitorTests.swift b/ios/PacketTunnelCoreTests/TunnelMonitorTests.swift
index 88b0c9c72c..7f9c16e99f 100644
--- a/ios/PacketTunnelCoreTests/TunnelMonitorTests.swift
+++ b/ios/PacketTunnelCoreTests/TunnelMonitorTests.swift
@@ -77,10 +77,10 @@ final class TunnelMonitorTests: XCTestCase {
         Task {
             for await event in await tunnelMonitor.eventStream {
                 switch event {
-                case .connectionEstablished:
-                    break
                 case .connectionLost:
                     connectionLostExpectation.fulfill()
+                default:
+                    break
                 }
             }
         }