[CHNL-16564] update script and message handling

Sources/KlaviyoUI/KlaviyoWebView/Development Assets/JSTestWebViewModel.swift

@@ -11,23 +11,34 @@ import Foundation
 import WebKit
 
 class JSTestWebViewModel: KlaviyoWebViewModeling {
-    let url: URL
-    let loadScripts: [String: WKUserScript]?
+    private enum MessageHandler: String, CaseIterable {
+        case toggleMessageHandler
+        case closeMessageHandler
+    }
+
     weak var delegate: KlaviyoWebViewDelegate?
 
+    let url: URL
+    var loadScripts: Set<WKUserScript>? = JSTestWebViewModel.initializeLoadScripts()
+    var messageHandlers: Set<String>? = Set(MessageHandler.allCases.map(\.rawValue))
+
     public let (navEventStream, navEventContinuation) = AsyncStream.makeStream(of: WKNavigationEvent.self)
 
     init(url: URL) {
         self.url = url
-        loadScripts = JSTestWebViewModel.initializeLoadScripts()
     }
 
-    private static func initializeLoadScripts() -> [String: WKUserScript] {
-        var scripts: [String: WKUserScript] = [:]
+    private static func initializeLoadScripts() -> Set<WKUserScript> {
+        var scripts = Set<WKUserScript>()
 
         if let toggleHandlerScript = try? ResourceLoader.getResourceContents(path: "toggleHandler", type: "js") {
             let script = WKUserScript(source: toggleHandlerScript, injectionTime: .atDocumentEnd, forMainFrameOnly: true)
-            scripts["toggleMessageHandler"] = script
+            scripts.insert(script)
+        }
+
+        if let closeHandlerScript = try? ResourceLoader.getResourceContents(path: "closeHandler", type: "js") {
+            let script = WKUserScript(source: closeHandlerScript, injectionTime: .atDocumentEnd, forMainFrameOnly: true)
+            scripts.insert(script)
         }
 
         return scripts
@@ -36,7 +47,13 @@ class JSTestWebViewModel: KlaviyoWebViewModeling {
     // MARK: handle WKWebView events
 
     func handleScriptMessage(_ message: WKScriptMessage) {
-        if message.name == "toggleMessageHandler" {
+        guard let handler = MessageHandler(rawValue: message.name) else {
+            // script message has no handler
+            return
+        }
+
+        switch handler {
+        case .toggleMessageHandler:
             guard let dict = message.body as? [String: AnyObject] else {
                 return
             }
@@ -61,6 +78,10 @@ class JSTestWebViewModel: KlaviyoWebViewModeling {
                     print("Javascript evaluation failed; message: \(error.localizedDescription)")
                 }
             }
+        case .closeMessageHandler:
+            Task {
+                await delegate?.dismiss()
+            }
         }
     }
 }

Sources/KlaviyoUI/KlaviyoWebView/KlaviyoWebViewController.swift

@@ -63,9 +63,8 @@ class KlaviyoWebViewController: UIViewController, WKUIDelegate, KlaviyoWebViewDe
     override func viewDidDisappear(_ animated: Bool) {
         super.viewDidDisappear(animated)
 
-        let scriptNames = viewModel.loadScripts?.keys.compactMap { $0 } ?? []
-        for scriptName in scriptNames {
-            webView.configuration.userContentController.removeScriptMessageHandler(forName: scriptName)
+        viewModel.messageHandlers?.forEach {
+            webView.configuration.userContentController.removeScriptMessageHandler(forName: $0)
         }
     }
 
@@ -90,11 +89,12 @@ class KlaviyoWebViewController: UIViewController, WKUIDelegate, KlaviyoWebViewDe
 
     /// Configures the scripts to be injected into the website when the website loads.
     private func configureLoadScripts() {
-        guard let scriptsDict = viewModel.loadScripts else { return }
+        viewModel.loadScripts?.forEach {
+            webView.configuration.userContentController.addUserScript($0)
+        }
 
-        for (name, script) in scriptsDict {
-            webView.configuration.userContentController.addUserScript(script)
-            webView.configuration.userContentController.add(self, name: name)
+        viewModel.messageHandlers?.forEach {
+            webView.configuration.userContentController.add(self, name: $0)
         }
     }
 

Sources/KlaviyoUI/KlaviyoWebView/KlaviyoWebViewModel.swift

@@ -23,23 +23,29 @@ public protocol KlaviyoWebViewDelegate: AnyObject {
 
 @_spi(KlaviyoPrivate)
 public class KlaviyoWebViewModel: KlaviyoWebViewModeling {
-    public let url: URL
-    public let loadScripts: [String: WKUserScript]?
+    private enum MessageHandler: String, CaseIterable {
+        case closeHandler
+    }
+
     public weak var delegate: KlaviyoWebViewDelegate?
 
+    public let url: URL
+    public let loadScripts: Set<WKUserScript>? = KlaviyoWebViewModel.initializeLoadScripts()
+    public var messageHandlers: Set<String>? = Set(MessageHandler.allCases.map(\.rawValue))
+
     public let (navEventStream, navEventContinuation) = AsyncStream.makeStream(of: WKNavigationEvent.self)
 
     public init(url: URL) {
         self.url = url
-        loadScripts = KlaviyoWebViewModel.initializeLoadScripts()
+//        loadScripts = KlaviyoWebViewModel.initializeLoadScripts()
     }
 
-    private static func initializeLoadScripts() -> [String: WKUserScript] {
-        var scripts: [String: WKUserScript] = [:]
+    private static func initializeLoadScripts() -> Set<WKUserScript> {
+        var scripts = Set<WKUserScript>()
 
         if let closeHandlerScript = try? ResourceLoader.getResourceContents(path: "closeHandler", type: "js") {
             let script = WKUserScript(source: closeHandlerScript, injectionTime: .atDocumentEnd, forMainFrameOnly: true)
-            scripts["closeHandler"] = script
+            scripts.insert(script)
         }
 
         return scripts
@@ -48,10 +54,13 @@ public class KlaviyoWebViewModel: KlaviyoWebViewModeling {
     // MARK: handle WKWebView events
 
     public func handleScriptMessage(_ message: WKScriptMessage) {
-        if message.name == "closeHandler" {
-            // TODO: handle close button tap
-            print("user tapped close button")
+        guard let handler = MessageHandler(rawValue: message.name) else {
+            // script message has no handler
+            return
+        }
 
+        switch handler {
+        case .closeHandler:
             Task {
                 await delegate?.dismiss()
             }

Sources/KlaviyoUI/KlaviyoWebView/KlaviyoWebViewModeling.swift

@@ -14,8 +14,10 @@ public protocol KlaviyoWebViewModeling: AnyObject {
     var url: URL { get }
     var delegate: KlaviyoWebViewDelegate? { get set }
 
-    /// Scripts to be injected into the ``WKWebView`` when the website loads.
-    var loadScripts: [String: WKUserScript]? { get }
+    /// Scripts & message handlers to be injected into the ``WKWebView`` when the website loads.
+    var loadScripts: Set<WKUserScript>? { get }
+    var messageHandlers: Set<String>? { get }
+
     var navEventStream: AsyncStream<WKNavigationEvent> { get }
     var navEventContinuation: AsyncStream<WKNavigationEvent>.Continuation { get }