fix preload timeout

Sources/KlaviyoUI/InAppForms/IAFWebViewModel.swift

@@ -24,7 +24,7 @@ class IAFWebViewModel: KlaviyoWebViewModeling {
     var messageHandlers: Set<String>? = Set(MessageHandler.allCases.map(\.rawValue))
 
     public let (navEventStream, navEventContinuation) = AsyncStream.makeStream(of: WKNavigationEvent.self)
-    private var formWillAppearContinuation: CheckedContinuation<Void, Never>?
+    private let (formWillAppearStream, formWillAppearContinuation) = AsyncStream.makeStream(of: Void.self)
 
     init(url: URL) {
         self.url = url
@@ -50,28 +50,47 @@ class IAFWebViewModel: KlaviyoWebViewModeling {
 
         do {
             try await withThrowingTaskGroup(of: Void.self) { group in
+                defer {
+                    formWillAppearContinuation.finish()
+                    group.cancelAll()
+                }
+
                 group.addTask {
                     try await Task.sleep(nanoseconds: timeout)
                     throw PreloadError.timeout
                 }
 
                 group.addTask { [weak self] in
-                    await withCheckedContinuation { continuation in
-                        self?.formWillAppearContinuation = continuation
-                    }
+                    guard let self else { return }
+
+                    var iterator = self.formWillAppearStream.makeAsyncIterator()
+                    await iterator.next()
                 }
 
-                if let _ = try await group.next() {
-                    // when the navigation task returns, we want to
-                    // cancel both the timeout task and the navigation task
-                    group.cancelAll()
+                group.addTask { [weak self] in
+                    guard let self else { return }
+                    for await event in self.navEventStream {
+                        if case .didFailNavigation = event {
+                            throw PreloadError.navigationFailed
+                        }
+                    }
                 }
+
+                try await group.next()
             }
-        } catch PreloadError.timeout {
-            if #available(iOS 14.0, *) {
-                Logger.webViewLogger.warning("Loading time exceeded specified timeout of \(Float(timeout / 1_000_000_000), format: .fixed(precision: 1)) seconds.")
+        } catch let error as PreloadError {
+            switch error {
+            case .timeout:
+                if #available(iOS 14.0, *) {
+                    Logger.webViewLogger.warning("Loading time exceeded specified timeout of \(Float(timeout / 1_000_000_000), format: .fixed(precision: 1)) seconds.")
+                }
+                throw error
+            case .navigationFailed:
+                if #available(iOS 14.0, *) {
+                    Logger.webViewLogger.warning("Navigation failed: \(error)")
+                }
+                throw error
             }
-            throw PreloadError.timeout
         } catch {
             if #available(iOS 14.0, *) {
                 Logger.webViewLogger.warning("Error preloading URL: \(error)")
@@ -110,8 +129,8 @@ class IAFWebViewModel: KlaviyoWebViewModeling {
             // TODO: handle formsDataLoaded
             ()
         case .formWillAppear:
-            formWillAppearContinuation?.resume()
-            formWillAppearContinuation = nil
+            formWillAppearContinuation.yield()
+            formWillAppearContinuation.finish()
         case let .trackAggregateEvent(data):
             KlaviyoInternal.create(aggregateEvent: data)
         case let .trackProfileEvent(data):

Tests/KlaviyoUITests/IAFWebViewModelPreloadingTests.swift

@@ -0,0 +1,195 @@
+//
+//  IAFWebViewModelTests.swift
+//  klaviyo-swift-sdk
+//
+//  Created by Andrew Balmer on 2/6/25.
+//
+
+@testable @_spi(KlaviyoPrivate) import KlaviyoUI
+import KlaviyoCore
+import WebKit
+import XCTest
+
+class MockIAFWebViewDelegate: NSObject, KlaviyoWebViewDelegate {
+    enum PreloadResult {
+        case formWillAppear(delay: UInt64)
+        case didFailNavigation(delay: UInt64)
+        case none
+    }
+
+    let viewModel: IAFWebViewModel
+
+    var preloadResult: PreloadResult?
+    var preloadUrlCalled = false
+    var evaluateJavaScriptCalled = false
+
+    init(viewModel: IAFWebViewModel) {
+        self.viewModel = viewModel
+    }
+
+    func preloadUrl() {
+        viewModel.handleNavigationEvent(.didCommitNavigation)
+        preloadUrlCalled = true
+
+        Task {
+            if let result = preloadResult {
+                switch result {
+                case let .formWillAppear(delay):
+                    try? await Task.sleep(nanoseconds: delay)
+
+                    let scriptMessage = MockWKScriptMessage(
+                        name: "KlaviyoNativeBridge",
+                        body: """
+                        {
+                          "type": "formWillAppear",
+                          "data": {
+                            "formId": "abc123"
+                          }
+                        }
+                        """)
+
+                    viewModel.handleScriptMessage(scriptMessage)
+
+                case let .didFailNavigation(delay):
+                    try? await Task.sleep(nanoseconds: delay)
+                    viewModel.handleNavigationEvent(.didFailNavigation)
+
+                case .none:
+                    // don't do anything
+                    return
+                }
+            }
+        }
+    }
+
+    func evaluateJavaScript(_ script: String) async throws -> Any {
+        evaluateJavaScriptCalled = true
+        return true
+    }
+
+    func dismiss() {}
+}
+
+class MockWKScriptMessage: WKScriptMessage {
+    private let mockName: String
+    private let mockBody: Any
+
+    init(name: String, body: Any) {
+        mockName = name
+        mockBody = body
+        super.init() // Calling the superclass initializer
+    }
+
+    override var name: String {
+        mockName
+    }
+
+    override var body: Any {
+        mockBody
+    }
+}
+
+final class IAFWebViewModelPreloadingTests: XCTestCase {
+    // MARK: - setup
+
+    var viewModel: IAFWebViewModel!
+    var delegate: MockIAFWebViewDelegate!
+
+    override func setUp() {
+        super.setUp()
+
+        viewModel = IAFWebViewModel(url: URL(string: "https://example.com")!)
+        delegate = MockIAFWebViewDelegate(viewModel: viewModel)
+        viewModel.delegate = delegate
+    }
+
+    override func tearDown() {
+        viewModel = nil
+        delegate = nil
+
+        super.tearDown()
+    }
+
+    // MARK: - tests
+
+    /// Tests scenario in which a `formWillAppear` event is emitted before the timeout is reached.
+    func testPreloadWebsiteSuccess() async throws {
+        // Given
+        delegate.preloadResult = .formWillAppear(delay: 100_000_000) // 0.1 second in nanoseconds
+        let expectation = XCTestExpectation(description: "Preloading website succeeds")
+
+        // When
+        do {
+            try await viewModel.preloadWebsite(timeout: 1_000_000_000) // 1 second in nanoseconds
+            expectation.fulfill()
+        } catch {
+            XCTFail("Expected success, but got error: \(error)")
+        }
+
+        // Then
+        XCTAssertTrue(delegate.preloadUrlCalled, "preloadUrl should be called on delegate")
+        await fulfillment(of: [expectation], timeout: 2.0)
+    }
+
+    /// Tests scenario in which the timeout is reached before the `formWillAppear` event is emitted.
+    func testPreloadWebsiteTimeout() async {
+        // Given
+        delegate.preloadResult = .formWillAppear(delay: 1_000_000_000) // 1 second in nanoseconds
+        let expectation = XCTestExpectation(description: "Preloading website times out")
+
+        // When
+        do {
+            try await viewModel.preloadWebsite(timeout: 100_000_000) // 0.1 second in nanoseconds
+            XCTFail("Expected timeout error, but succeeded")
+        } catch PreloadError.timeout {
+            expectation.fulfill()
+        } catch {
+            XCTFail("Expected timeout error, but got: \(error)")
+        }
+
+        // Then
+        XCTAssertTrue(delegate.preloadUrlCalled, "preloadUrl should be called on delegate")
+        await fulfillment(of: [expectation], timeout: 2.0)
+    }
+
+    /// Tests scenario in which the delegate does nothing and emits no events after `preloadUrl()` is called.
+    func testPreloadWebsiteNoActionTimeout() async {
+        // Given
+        delegate.preloadResult = MockIAFWebViewDelegate.PreloadResult.none
+        let expectation = XCTestExpectation(description: "Preloading website times out")
+
+        // When
+        do {
+            try await viewModel.preloadWebsite(timeout: 100_000_000) // 0.1 second in nanoseconds
+            XCTFail("Expected timeout error, but succeeded")
+        } catch PreloadError.timeout {
+            expectation.fulfill()
+        } catch {
+            XCTFail("Expected timeout error, but got: \(error)")
+        }
+
+        // Then
+        XCTAssertTrue(delegate.preloadUrlCalled, "preloadUrl should be called on delegate")
+        await fulfillment(of: [expectation], timeout: 2.0)
+    }
+
+    func testPreloadWebsiteNavigationFailed() async {
+        // Given
+        delegate.preloadResult = .didFailNavigation(delay: 100_000_000) // 0.1 second in nanoseconds
+        let expectation = XCTestExpectation(description: "Preloading website fails")
+
+        // When
+        do {
+            try await viewModel.preloadWebsite(timeout: 1_000_000_000) // 1 second in nanoseconds
+            XCTFail("Expected navigation failed error, but succeeded")
+        } catch PreloadError.navigationFailed {
+            expectation.fulfill()
+        } catch {
+            XCTFail("Expected navigation failed error, but got: \(error)")
+        }
+
+        // Then
+        XCTAssertTrue(delegate.preloadUrlCalled, "preloadUrl should be called on delegate")
+        await fulfillment(of: [expectation], timeout: 2.0)
+    }
+}

Tests/KlaviyoUITests/IAFWebViewModelTests.swift

@@ -40,7 +40,7 @@ final class IAFWebViewModelTests: XCTestCase {
         super.tearDown()
     }
 
-    // MARK: - tests
+    // MARK: - html injection tests
 
     func testInjectSdkNameAttribute() async throws {
         // This test has been flaky when running on CI. It seems to have something to do with instability when