Fixing WebViewCaptureService proxy

Sources/EmbraceCore/Capture/WebView/WKNavigationDelegateProxy.swift

@@ -12,6 +12,11 @@ class WKNavigationDelegateProxy: NSObject {
     // callback triggered the webview loads an url or errors
     var callback: ((URL?, Int?) -> Void)?
 
+    init(originalDelegate: WKNavigationDelegate? = nil, callback: ((URL?, Int?) -> Void)? = nil) {
+        self.originalDelegate = originalDelegate
+        self.callback = callback
+    }
+
     override func responds(to aSelector: Selector!) -> Bool {
         if super.responds(to: aSelector) {
             return true

Sources/EmbraceCore/Capture/WebView/WKWebView+Embrace.swift

@@ -0,0 +1,32 @@
+//
+//  Copyright © 2024 Embrace Mobile, Inc. All rights reserved.
+//
+
+#if canImport(WebKit)
+import Foundation
+import WebKit
+
+extension WKWebView {
+    private struct AssociatedKeys {
+        static var embraceProxy: Int = 0
+    }
+
+    var emb_proxy: WKNavigationDelegateProxy? {
+        get {
+            if let value = objc_getAssociatedObject(self, &AssociatedKeys.embraceProxy) as? WKNavigationDelegateProxy {
+                return value as WKNavigationDelegateProxy
+            }
+
+            return nil
+        }
+
+        set {
+            objc_setAssociatedObject(self,
+                                     &AssociatedKeys.embraceProxy,
+                                     newValue,
+                                     .OBJC_ASSOCIATION_RETAIN)
+        }
+    }
+}
+
+#endif

Sources/EmbraceCore/Capture/WebView/WebViewCaptureService.swift

@@ -18,7 +18,6 @@ public final class WebViewCaptureService: CaptureService {
     @objc public let options: WebViewCaptureService.Options
     private let lock: NSLocking
     private var swizzlers: [any Swizzlable] = []
-    var proxy: WKNavigationDelegateProxy
 
     @objc public convenience init(options: WebViewCaptureService.Options) {
         self.init(options: options, lock: NSLock())
@@ -34,13 +33,8 @@ public final class WebViewCaptureService: CaptureService {
     ) {
         self.options = options
         self.lock = lock
-        self.proxy = WKNavigationDelegateProxy()
 
         super.init()
-
-        proxy.callback = { [weak self] url, statusCode in
-            self?.createEvent(url: url, statusCode: statusCode)
-        }
     }
 
     public override func onInstall() {
@@ -65,14 +59,17 @@ public final class WebViewCaptureService: CaptureService {
     }
 
     private func initializeSwizzlers() {
-        swizzlers.append(WKWebViewSetNavigationDelegateSwizzler(proxy: proxy))
+        swizzlers.append(WKWebViewSetNavigationDelegateSwizzler(delegate: self))
         swizzlers.append(WKWebViewLoadRequestSwizzler())
         swizzlers.append(WKWebViewLoadHTMLStringSwizzler())
         swizzlers.append(WKWebViewLoadFileURLSwizzler())
         swizzlers.append(WKWebViewLoadDataSwizzler())
     }
+}
+
+extension WebViewCaptureService: WebViewSwizzlerDelegate {
 
-    private func createEvent(url: URL?, statusCode: Int?) {
+    func didLoad(url: URL?, statusCode: Int?) {
         guard let url = url else {
             return
         }
@@ -116,18 +113,23 @@ struct WKWebViewSetNavigationDelegateSwizzler: Swizzlable {
     typealias BlockImplementationType = @convention(block) (WKWebView, WKNavigationDelegate) -> Void
     static var selector: Selector = #selector(setter: WKWebView.navigationDelegate)
     var baseClass: AnyClass
-    let proxy: WKNavigationDelegateProxy
+    weak var delegate: WebViewSwizzlerDelegate?
 
-    init(proxy: WKNavigationDelegateProxy, baseClass: AnyClass = WKWebView.self) {
+    init(delegate: WebViewSwizzlerDelegate?, baseClass: AnyClass = WKWebView.self) {
         self.baseClass = baseClass
-        self.proxy = proxy
+        self.delegate = delegate
     }
 
     func install() throws {
         try swizzleInstanceMethod { originalImplementation -> BlockImplementationType in
             return { webView, delegate in
                 if !(webView.navigationDelegate is WKNavigationDelegateProxy) {
-                    proxy.originalDelegate = delegate
+                    let proxy = WKNavigationDelegateProxy(
+                        originalDelegate: delegate) { url, statusCode in
+                            self.delegate?.didLoad(url: url, statusCode: statusCode)
+                        }
+                    webView.emb_proxy = proxy
+
                     originalImplementation(webView, Self.selector, proxy)
                 } else {
                     originalImplementation(webView, Self.selector, delegate)
@@ -230,4 +232,8 @@ struct WKWebViewLoadDataSwizzler: Swizzlable {
         }
     }
 }
+
+protocol WebViewSwizzlerDelegate: AnyObject {
+    func didLoad(url: URL?, statusCode: Int?)
+}
 #endif

Tests/EmbraceCoreTests/Capture/UX/View/UIViewControllerHandlerTests.swift

@@ -172,7 +172,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         validateViewDidAppearSpans(vc: vc, parentName: parentName)
 
         // then all the spans are created and ended at the right times
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             return parent != nil && self.cacheIsEmpty()
         }
@@ -190,7 +190,7 @@ class UIViewControllerHandlerTests: XCTestCase {
 
         handler.onViewDidDisappear(vc)
 
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             return parent != nil && parent!.status.isError == true && self.cacheIsEmpty()
         }
@@ -208,7 +208,7 @@ class UIViewControllerHandlerTests: XCTestCase {
 
         handler.appDidEnterBackground()
 
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             return parent != nil && parent!.status.isError == true && self.cacheIsEmpty()
         }
@@ -238,7 +238,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewBecameInteractive(vc)
 
         // then the spans are ended
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             let uiReady = self.otel.spanProcessor.endedSpans.first(where: { $0.name == "ui-ready"})
 
@@ -260,7 +260,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         validateViewDidAppearSpans(vc: vc, parentName: parentName)
 
         // then the spans are ended
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             let uiReady = self.otel.spanProcessor.endedSpans.first(where: { $0.name == "ui-ready"})
 
@@ -280,7 +280,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewDidDisappear(vc)
 
         // then the spans are ended
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             return parent != nil && parent!.status.isError == true && self.cacheIsEmpty()
         }
@@ -299,7 +299,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.appDidEnterBackground()
 
         // then the spans are ended
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.endedSpans.first(where: { $0.name.contains(parentName) })
             return parent != nil && parent!.status.isError == true && self.cacheIsEmpty()
         }
@@ -310,7 +310,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewDidLoadStart(vc)
 
         // then spans are created
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.startedSpans.first(where: { $0.name.contains(parentName) })
             let child = self.otel.spanProcessor.startedSpans.first(where: { $0.name == "emb-view-did-load"})
 
@@ -321,7 +321,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewDidLoadEnd(vc)
 
         // then the view did load span is ended
-        wait {
+        wait(timeout: .longTimeout) {
             let span = self.otel.spanProcessor.startedSpans.first(where: { $0.name == "emb-view-did-load"})
             return span != nil && self.handler.viewDidLoadSpans.isEmpty
         }
@@ -332,7 +332,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewWillAppearStart(vc)
 
         // then a child span is created
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.startedSpans.first(where: { $0.name.contains(parentName) })
             let child = self.otel.spanProcessor.startedSpans.first(where: { $0.name == "emb-view-will-appear"})
 
@@ -343,7 +343,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewWillAppearEnd(vc)
 
         // then the view will appear span is ended
-        wait {
+        wait(timeout: .longTimeout) {
             let span = self.otel.spanProcessor.endedSpans.first(where: { $0.name == "emb-view-will-appear"})
             return span != nil && self.handler.viewWillAppearSpans.isEmpty
         }
@@ -354,7 +354,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewDidAppearStart(vc)
 
         // then a child span is created
-        wait {
+        wait(timeout: .longTimeout) {
             let parent = self.otel.spanProcessor.startedSpans.first(where: { $0.name.contains(parentName) })
             let child = self.otel.spanProcessor.startedSpans.first(where: { $0.name == "emb-view-did-appear"})
 
@@ -365,7 +365,7 @@ class UIViewControllerHandlerTests: XCTestCase {
         handler.onViewDidAppearEnd(vc)
 
         // then the view did appear span is ended
-        wait {
+        wait(timeout: .longTimeout) {
             let span = self.otel.spanProcessor.endedSpans.first(where: { $0.name == "emb-view-did-appear"})
             return span != nil && self.handler.viewDidAppearSpans.isEmpty
         }

Tests/EmbraceCoreTests/Capture/WebView/WebViewCaptureServiceTests.swift

@@ -30,8 +30,8 @@ class WebViewCaptureServiceTests: XCTestCase {
 
         // then a proxy delegate is correctly set
         XCTAssert(webView.navigationDelegate!.isKind(of: WKNavigationDelegateProxy.self))
-        XCTAssertNotNil(service.proxy.originalDelegate)
-        XCTAssert(service.proxy.originalDelegate!.isKind(of: MockWKNavigationDelegate.self))
+        XCTAssertNotNil(webView.emb_proxy!.originalDelegate)
+        XCTAssert(webView.emb_proxy!.originalDelegate!.isKind(of: MockWKNavigationDelegate.self))
     }
 
     func test_setNavigationDelegate_ShouldntGenerateRecursion() throws {
@@ -49,19 +49,9 @@ class WebViewCaptureServiceTests: XCTestCase {
     }
 
     func test_spanEvent() {
-        // given a webview
-        let webView = WKWebView()
-
-        // when loading an URL
+        // when a url is loaded
         let url = URL(string: "https://www.google.com/")!
-        let request = URLRequest(url: url)
-        webView.load(request)
-        wait(delay: .longTimeout)
-
-        // manually call this for test to pass on CI
-        service.proxy.webView(webView, decidePolicyFor: response) { _ in
-
-        }
+        service.didLoad(url: url, statusCode: nil)
 
         // then a span event is created
         XCTAssert(otel.events.count > 0)
@@ -73,18 +63,9 @@ class WebViewCaptureServiceTests: XCTestCase {
     }
 
     func test_spanEvent_withError() {
-        // given a webview
-        let webView = WKWebView()
-
-        // when loading an URL
+        // when a url is loaded with error
         let url = URL(string: "https://www.google.com/")!
-        let request = URLRequest(url: url)
-        webView.load(request)
-        wait(delay: .longTimeout)
-
-        // when an error occurs
-        let error = NSError(domain: TestConstants.domain, code: 123)
-        service.proxy.webView(webView, didFail: navigation, withError: error)
+        service.didLoad(url: url, statusCode: 123)
 
         // then a span event is created with an error code
         XCTAssert(otel.events.count > 0)