From a41ea1ed0958246db0d7aedbd4ed46fa3e82f10f Mon Sep 17 00:00:00 2001
From: Asaf Korem <asaf.korem@gmail.com>
Date: Sat, 19 Oct 2024 14:24:27 +0300
Subject: [PATCH] fix(iOS): wait for webview to load before injecting JS to
 extract html.

---
 .../Utilities/ViewHierarchyGenerator.swift    | 58 +++++++++++++++++--
 1 file changed, 54 insertions(+), 4 deletions(-)

diff --git a/detox/ios/Detox/Utilities/ViewHierarchyGenerator.swift b/detox/ios/Detox/Utilities/ViewHierarchyGenerator.swift
index dbdee726cb..9b664ac930 100644
--- a/detox/ios/Detox/Utilities/ViewHierarchyGenerator.swift
+++ b/detox/ios/Detox/Utilities/ViewHierarchyGenerator.swift
@@ -100,15 +100,23 @@ struct ViewHierarchyGenerator {
 
     @MainActor
     private static func getHtmlFromWebView(_ webView: WKWebView) async -> String {
-        await withCheckedContinuation { continuation in
+        let handler = WebViewHandler()
+
+        do {
+            try await handler.waitForPageLoad(webView)
+        } catch {
+            print("Error waiting for page to load: \(error)")
+            return "<!-- Error loading page -->"
+        }
+
+        return await withCheckedContinuation { continuation in
             webView.evaluateJavaScript(GET_HTML_SCRIPT) { result, error in
                 if let html = result as? String {
                     continuation.resume(returning: html)
                 } else if let error = error {
-                    print("Error extracting HTML: \(error)")
-                    continuation.resume(returning: "<!-- Error loading HTML -->")
+                    continuation.resume(returning: "<!-- Error loading HTML: \(error) -->")
                 } else {
-                    continuation.resume(returning: "html is empty")
+                    continuation.resume(returning: "<!--- HTML is empty -->")
                 }
             }
         }
@@ -171,3 +179,45 @@ struct ViewHierarchyGenerator {
             .joined()
     }
 }
+
+@MainActor
+class WebViewHandler: NSObject, WKNavigationDelegate {
+    private var loadCompletion: ((Result<Void, Error>) -> Void)?
+
+    enum WebViewError: Error {
+        case timeout
+    }
+
+    func waitForPageLoad(_ webView: WKWebView, timeoutAfter seconds: TimeInterval = 10) async throws {
+        webView.navigationDelegate = self
+
+        if !webView.isLoading {
+            return
+        }
+
+        return try await withCheckedThrowingContinuation { continuation in
+            let task = Task {
+                try await Task.sleep(nanoseconds: UInt64(seconds * 1_000_000_000))
+                loadCompletion?(.failure(WebViewError.timeout))
+            }
+
+            self.loadCompletion = { result in
+                task.cancel()
+                self.loadCompletion = nil
+                continuation.resume(with: result)
+            }
+        }
+    }
+
+    nonisolated func webView(_ webView: WKWebView, didFinish navigation: WKNavigation!) {
+        Task { @MainActor in
+            loadCompletion?(.success(()))
+        }
+    }
+
+    nonisolated func webView(_ webView: WKWebView, didFail navigation: WKNavigation!, withError error: Error) {
+        Task { @MainActor in
+            loadCompletion?(.failure(error))
+        }
+    }
+}