RUMM-1931 PR comments addressed

DataDog · Feb 8, 2022 · deb2ddd · deb2ddd
1 parent 773f394
commit deb2ddd
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 10 deletions.
diff --git a/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift b/Datadog/Example/Scenarios/WebView/WebViewTrackingFixtureViewController.swift
@@ -27,15 +27,20 @@ class ShopistWebviewViewController: UIViewController {
     override func viewDidLoad() {
         super.viewDidLoad()
 
-        let controller = WKUserContentController()
-        controller.trackDatadogEvents(in: ["shopist.io"])
-        let config = WKWebViewConfiguration()
-        config.userContentController = controller
-
-        webView = WKWebView(frame: UIScreen.main.bounds, configuration: config)
+        webView = WKWebView(frame: UIScreen.main.bounds, configuration: WKWebViewConfiguration())
         view.addSubview(webView)
     }
 
+    override func viewWillAppear(_ animated: Bool) {
+        super.viewWillAppear(animated)
+        webView.configuration.userContentController.trackDatadogEvents(in: ["shopist.io"])
+    }
+
+    override func viewWillDisappear(_ animated: Bool) {
+        super.viewWillDisappear(animated)
+        webView.configuration.userContentController.stopTrackingDatadogEvents()
+    }
+
     override func viewDidAppear(_ animated: Bool) {
         super.viewDidAppear(animated)
         webView.load(request)

diff --git a/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift b/Sources/Datadog/FeaturesIntegration/WebView/WKUserContentController+Datadog.swift
@@ -30,7 +30,7 @@ public extension WKUserContentController {
     /// Disables Datadog iOS SDK and Datadog Browser SDK integration.
     ///
     /// Removes Datadog's ScriptMessageHandler and UserScript from the caller.
-    /// _NOTE:_ This method **must** be called when the webview can be deinitialized.
+    /// - Note: This method **must** be called when the webview can be deinitialized.
     func stopTrackingDatadogEvents() {
         removeScriptMessageHandler(forName: DatadogMessageHandler.name)
 

diff --git a/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift b/Tests/DatadogTests/Datadog/RUM/WebView/WKUserContentController+DatadogTests.swift
@@ -24,9 +24,7 @@ final class DDUserContentController: WKUserContentController {
 }
 
 final class MockMessageHandler: NSObject, WKScriptMessageHandler {
-    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) {
-        return
-    }
+    func userContentController(_ userContentController: WKUserContentController, didReceive message: WKScriptMessage) { }
 }
 
 final class MockScriptMessage: WKScriptMessage {