Fix traffic generation

4 files changed, 75 insertions, 32 deletions

diff --git a/ios/MullvadVPNUITests/Base/BaseUITestCase.swift b/ios/MullvadVPNUITests/Base/BaseUITestCase.swift
index d7443bea8c..7b978a61f0 100644
--- a/ios/MullvadVPNUITests/Base/BaseUITestCase.swift
+++ b/ios/MullvadVPNUITests/Base/BaseUITestCase.swift
@@ -97,7 +97,7 @@ class BaseUITestCase: XCTestCase {
 
     /// Create temporary account without time. Will be created using partner API if token is configured, else falling back to app API
     func createTemporaryAccountWithoutTime() -> String {
-        if let partnerApiToken {
+        if partnerApiToken != nil {
             let partnerAPIClient = PartnerAPIClient()
             return partnerAPIClient.createAccount()
         } else {
@@ -401,4 +401,5 @@ class BaseUITestCase: XCTestCase {
             XCTFail("Failed to find 'Delete'")
         }
     }
+    // swiftlint:disable:next file_length
 }
diff --git a/ios/MullvadVPNUITests/LeakTests.swift b/ios/MullvadVPNUITests/LeakTests.swift
index a6daffcff3..1ad640f9a3 100644
--- a/ios/MullvadVPNUITests/LeakTests.swift
+++ b/ios/MullvadVPNUITests/LeakTests.swift
@@ -19,15 +19,15 @@ class LeakTests: LoggedInWithTimeUITestCase {
         let targetIPAddress = Networking.getAlwaysReachableIPAddress()
         startPacketCapture()
         let trafficGenerator = TrafficGenerator(destinationHost: targetIPAddress, port: 80)
-        trafficGenerator.startGeneratingUDPTraffic(interval: 30.0)
+        trafficGenerator.startGeneratingUDPTraffic(interval: 1.0)
 
         TunnelControlPage(app)
-            .tapSecureConnectionButton()
+            .tapConnectButton()
 
         allowAddVPNConfigurationsIfAsked()
 
         TunnelControlPage(app)
-            .waitForSecureConnectionLabel()
+            .waitForConnectedLabel()
 
         // Keep the tunnel connection for a while
         Thread.sleep(forTimeInterval: 30.0)
@@ -51,33 +51,33 @@ class LeakTests: LoggedInWithTimeUITestCase {
         trafficGenerator.startGeneratingUDPTraffic(interval: 1.0)
 
         TunnelControlPage(app)
-            .tapSecureConnectionButton()
+            .tapConnectButton()
 
         allowAddVPNConfigurationsIfAsked()
 
         TunnelControlPage(app)
-            .waitForSecureConnectionLabel()
+            .waitForConnectedLabel()
 
-        Thread.sleep(forTimeInterval: 2.0)
+        RunLoop.current.run(until: .now + 2)
 
         TunnelControlPage(app)
             .tapDisconnectButton()
 
         // Give it some time to generate traffic outside of tunnel
-        Thread.sleep(forTimeInterval: 5.0)
+        RunLoop.current.run(until: .now + 5)
 
         TunnelControlPage(app)
-            .tapSecureConnectionButton()
+            .tapConnectButton()
 
         // Keep the tunnel connection for a while
-        Thread.sleep(forTimeInterval: 5.0)
+        RunLoop.current.run(until: .now + 5)
 
         app.launch()
         TunnelControlPage(app)
             .tapDisconnectButton()
 
         // Keep the capture open for a while
-        Thread.sleep(forTimeInterval: 15.0)
+        RunLoop.current.run(until: .now + 15)
         trafficGenerator.stopGeneratingUDPTraffic()
 
         var capturedStreams = stopPacketCapture()
diff --git a/ios/MullvadVPNUITests/Networking/PacketCapture.swift b/ios/MullvadVPNUITests/Networking/PacketCapture.swift
index dd80f06580..7c26d9cc59 100644
--- a/ios/MullvadVPNUITests/Networking/PacketCapture.swift
+++ b/ios/MullvadVPNUITests/Networking/PacketCapture.swift
@@ -10,7 +10,13 @@ import Foundation
 import XCTest
 
 struct PacketCaptureSession {
-    var identifier = UUID().uuidString
+    var identifier: String
+
+    init(identifier: String = UUID().uuidString) {
+        self.identifier = identifier
+
+        print("Current Packet Capture session identifier is: \(identifier)")
+    }
 }
 
 /// Represents a stream in packet capture
diff --git a/ios/MullvadVPNUITests/Networking/TrafficGenerator.swift b/ios/MullvadVPNUITests/Networking/TrafficGenerator.swift
index b339e32819..335a1b6455 100644
--- a/ios/MullvadVPNUITests/Networking/TrafficGenerator.swift
+++ b/ios/MullvadVPNUITests/Networking/TrafficGenerator.swift
@@ -12,63 +12,99 @@ import XCTest
 class TrafficGenerator {
     let destinationHost: String
     let port: Int
-    let connection: NWConnection
-    let dispatchQueue = DispatchQueue(label: "TrafficGeneratorDispatchQueue", qos: .unspecified)
-    var timer: DispatchSourceTimer
+    var connection: NWConnection
+    let dispatchQueue = DispatchQueue(label: "TrafficGeneratorDispatchQueue")
+    var sendDataTimer: DispatchSourceTimer
 
     init(destinationHost: String, port: Int) {
         self.destinationHost = destinationHost
         self.port = port
+
+        sendDataTimer = DispatchSource.makeTimerSource(queue: dispatchQueue)
+        let params = NWParameters.udp
         connection = NWConnection(
             host: NWEndpoint.Host(destinationHost),
             port: NWEndpoint.Port(integerLiteral: UInt16(port)),
-            using: .udp
+            using: params
         )
+        setupConnection()
+    }
+
+    func reconnect() {
+        print("Attempting to reconnect")
+        self.connection.forceCancel()
 
-        timer = DispatchSource.makeTimerSource(queue: dispatchQueue)
+        connection = recreateConnection()
+        self.setupConnection()
+    }
 
-        connect()
+    func recreateConnection() -> NWConnection {
+        let params = NWParameters.udp
+        return NWConnection(
+            host: NWEndpoint.Host(destinationHost),
+            port: NWEndpoint.Port(integerLiteral: UInt16(port)),
+            using: params
+        )
     }
 
-    func connect() {
+    func setupConnection() {
+        print("Setting up connection...")
         let doneAttemptingConnectExpecation = XCTestExpectation(description: "Done attemping to connect")
 
         connection.stateUpdateHandler = { state in
             switch state {
             case .ready:
                 print("Ready")
+                self.sendDataTimer.resume()
                 doneAttemptingConnectExpecation.fulfill()
             case let .failed(error):
                 print("Failed to connect: \(error)")
-                doneAttemptingConnectExpecation.fulfill()
+                self.sendDataTimer.cancel()
+                self.reconnect()
+            case .preparing:
+                print("Preparing connection...")
+            case .setup:
+                print("Setting upp connection...")
+            case let .waiting(error):
+                print("Waiting to connect: \(error)")
+            case .cancelled:
+                self.sendDataTimer.suspend()
+                print("Cancelled connection")
+                self.reconnect()
             default:
                 break
             }
         }
-
         connection.start(queue: dispatchQueue)
 
         XCTWaiter().wait(for: [doneAttemptingConnectExpecation], timeout: 10.0)
     }
 
     public func startGeneratingUDPTraffic(interval: TimeInterval) {
-        timer.schedule(deadline: .now(), repeating: interval)
+        sendDataTimer.schedule(deadline: .now(), repeating: interval)
 
-        timer.setEventHandler {
+        sendDataTimer.setEventHandler {
             let data = "dGhpcyBpcyBqdXN0IHNvbWUgZHVtbXkgZGF0YSB0aGlzIGlzIGp1c3Qgc29tZSBkdW".data(using: .utf8)
-            self.connection.send(content: data, completion: .contentProcessed { error in
-                if let error = error {
-                    print("Failed to send data: \(error)")
-                } else {
-                    print("Data sent")
-                }
-            })
+
+            print("Attempting to send data...")
+
+            if self.connection.state != .ready {
+                print("Not connected, won't send data")
+            } else {
+                self.connection.send(content: data, completion: .contentProcessed { error in
+                    if let error = error {
+                        print("Failed to send data: \(error)")
+                    } else {
+                        print("Data sent")
+                    }
+                })
+            }
         }
 
-        timer.activate()
+        sendDataTimer.activate()
     }
 
     public func stopGeneratingUDPTraffic() {
-        timer.cancel()
+        sendDataTimer.cancel()
     }
 }