From 2766f667246d3eae257a3a6234b6dc1ee2132b31 Mon Sep 17 00:00:00 2001
From: louiszawadzki <louis.zawadzki@datadoghq.com>
Date: Thu, 14 Sep 2023 11:08:15 +0200
Subject: [PATCH] Memoize webview message emitter

---
 .../ios/Sources/DdSdkImplementation.swift     |  9 ++++----
 packages/core/ios/Tests/DdSdkTests.swift      | 21 +++++++++++++++++++
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/packages/core/ios/Sources/DdSdkImplementation.swift b/packages/core/ios/Sources/DdSdkImplementation.swift
index b7b4821ef..d97445df5 100644
--- a/packages/core/ios/Sources/DdSdkImplementation.swift
+++ b/packages/core/ios/Sources/DdSdkImplementation.swift
@@ -27,7 +27,8 @@ public class DdSdkImplementation: NSObject {
     let mainDispatchQueue: DispatchQueueType
     let RUMMonitorProvider: () -> RUMMonitorProtocol
     let RUMMonitorInternalProvider: () -> RUMMonitorInternalProtocol?
-    
+    var webviewMessageEmitter: InternalExtension<WebViewTracking>.AbstractMessageEmitter?
+
     private let jsLongTaskThresholdInSeconds: TimeInterval = 0.1;
 
     @objc
@@ -101,6 +102,8 @@ public class DdSdkImplementation: NSObject {
         if sdkConfiguration.nativeCrashReportEnabled ?? false {
             CrashReporting.enable(in: core)
         }
+        
+        self.webviewMessageEmitter = WebViewTracking._internal.messageEmitter(in: core)
 
         overrideReactNativeTelemetry(rnConfiguration: sdkConfiguration, core: core)
     }
@@ -149,9 +152,7 @@ public class DdSdkImplementation: NSObject {
     @objc
     public func consumeWebviewEvent(message: NSString, resolve:RCTPromiseResolveBlock, reject:RCTPromiseRejectBlock) -> Void {
         do{
-            // TODO: memoize message emitter once core is initialized
-            let messageEmitter = WebViewTracking._internal.messageEmitter(in: CoreRegistry.default)
-            try messageEmitter.send(body: message)
+            try self.webviewMessageEmitter?.send(body: message)
         } catch {
             Datadog._internal.telemetry.error(id: "datadog_react_native:\(error.localizedDescription)", message: "The message being sent was:\(message)" as String, kind: "WebViewEventBridgeError" as String, stack: String(describing: error) as String)
         }
diff --git a/packages/core/ios/Tests/DdSdkTests.swift b/packages/core/ios/Tests/DdSdkTests.swift
index df4a4c840..557f0c18d 100644
--- a/packages/core/ios/Tests/DdSdkTests.swift
+++ b/packages/core/ios/Tests/DdSdkTests.swift
@@ -874,6 +874,22 @@ internal class DdSdkTests: XCTestCase {
 
         Datadog.internalFlushAndDeinitialize()
     }
+
+    func testConsumeWebviewEventBeforeInitialization() throws {
+        XCTAssertNoThrow(try DdSdkImplementation().consumeWebviewEvent(message: "TestMessage", resolve: mockResolve, reject: mockReject))
+    }
+
+    func testConsumeWebviewEvent() throws {
+        let sdk = DdSdkImplementation()
+        let configuration: DdSdkConfiguration = .mockAny()
+        let core = MockDatadogCore()
+
+        sdk.enableFeatures(sdkConfiguration: configuration, core: core)
+        
+        sdk.consumeWebviewEvent(message: "{\"eventType\":\"RUM\",\"event\":{\"blabla\":\"custom message\"}}", resolve: mockResolve, reject: mockReject)
+        
+        XCTAssertNotNil(core.baggages["browser-rum-event"])
+    }
 }
 
 private final class MockJSRefreshRateMonitor: RefreshRateMonitor {
@@ -998,10 +1014,15 @@ internal class MockDatadogCore: DatadogCoreProtocol {
             case .configuration(let configuration) = telemetry {
             self.configuration = configuration
         }
+        
+        if case .baggage(let key, let baggage) = message {
+            self.baggages[key] = baggage
+        }
     }
    
     private(set) var configuration: ConfigurationTelemetry?
     private(set) var features: [String: DatadogFeature] = [:]
+    private(set) var baggages: [String: Any] = [:]
 
     func register<T>(feature: T) throws where T : DatadogFeature {
         features[T.name] = feature