using cli changes for cross origin iframe

Author: yashmahamulkar-bs

percy/snapshot.py

@@ -1,7 +1,6 @@
 import os
 import platform
 import json
-import re
 from contextlib import contextmanager
 from functools import lru_cache
 from time import sleep
@@ -86,7 +85,7 @@ def fetch_percy_dom():
     response.raise_for_status()
     return response.text
 
-# pylint: disable=too-many-arguments, too-many-branches, too-many-locals
+# pylint: disable=too-many-arguments, too-many-branches, too-many-locals, too-many-positional-arguments
 def create_region(
     boundingBox=None,
     elementXpath=None,
@@ -175,10 +174,6 @@ def process_frame(driver, frame_element, options, percy_dom_script):
     }
 
 
-def _replace_iframe_with_srcdoc(match, srcdoc_value):
-    return f'{match.group(1)} srcdoc="{srcdoc_value}">'
-
-
 def _is_unsupported_iframe_src(frame_src):
     return (
         not frame_src or
@@ -193,55 +188,6 @@ def _get_origin(url):
     parsed = urlparse(url)
     return f"{parsed.scheme}://{parsed.netloc}"
 
-def stitch_cors_iframes(dom_snapshot, processed_frames):
-    html = dom_snapshot.get("html")
-    if not isinstance(html, str):
-        return dom_snapshot
-
-    combined_resources = []
-    if isinstance(dom_snapshot.get("resources"), list):
-        combined_resources = dom_snapshot["resources"][:]
-
-    seen_resource_urls = {
-        resource.get("url") for resource in combined_resources
-        if isinstance(resource, dict) and isinstance(resource.get("url"), str)
-    }
-
-    for frame in processed_frames:
-        iframe_data = frame.get("iframeData", {})
-        iframe_snapshot = frame.get("iframeSnapshot", {})
-        percy_element_id = iframe_data.get("percyElementId")
-        iframe_html = iframe_snapshot.get("html") if isinstance(iframe_snapshot, dict) else None
-
-        if not percy_element_id or not iframe_html:
-            continue
-
-        srcdoc_value = iframe_html.replace("&", "&").replace('"', """)
-        escaped_id = re.escape(percy_element_id)
-        iframe_regex = re.compile(
-            rf'(<iframe\b[^>]*?data-percy-element-id="{escaped_id}"[^>]*?)(/?>)',
-            re.DOTALL
-        )
-
-        html = iframe_regex.sub(
-            lambda match, srcdoc=srcdoc_value: _replace_iframe_with_srcdoc(match, srcdoc),
-            html,
-            count=1
-        )
-
-        resources = iframe_snapshot.get("resources") if isinstance(iframe_snapshot, dict) else None
-        if isinstance(resources, list):
-            for resource in resources:
-                if not isinstance(resource, dict):
-                    continue
-                url = resource.get("url")
-                if not isinstance(url, str) or url in seen_resource_urls:
-                    continue
-                combined_resources.append(resource)
-                seen_resource_urls.add(url)
-
-    return {**dom_snapshot, "html": html, "resources": combined_resources}
-
 def get_serialized_dom(driver, cookies, percy_dom_script=None, **kwargs):
     # 1. Serialize the main page first (this adds the data-percy-element-ids)
     dom_snapshot = driver.execute_script(f'return PercyDOM.serialize({json.dumps(kwargs)})')
@@ -270,15 +216,14 @@ def get_serialized_dom(driver, cookies, percy_dom_script=None, **kwargs):
                     processed_frames.append(result)
 
             if processed_frames:
-                dom_snapshot = stitch_cors_iframes(dom_snapshot, processed_frames)
+                dom_snapshot['corsIframes'] = processed_frames
     except Exception as e:
         log(f"Failed to process cross-origin iframes: {e}", "debug")
 
     dom_snapshot['cookies'] = cookies
     return dom_snapshot
 
 def get_responsive_widths(widths=None):
-    """Gets computed responsive widths from the Percy server for responsive snapshot capture."""
     if widths is None:
         widths = []
     try:
@@ -300,7 +245,6 @@ def get_responsive_widths(widths=None):
         msg = "Update Percy CLI to the latest version to use responsiveSnapshotCapture"
         raise Exception(msg) from e
 
-
 def _setup_resize_listener(driver):
     """Initializes the resize counter and attaches a named listener to avoid duplicates."""
     driver.execute_script("""

tests/test_snapshot.py

@@ -992,21 +992,16 @@ def test_process_frame_returns_none_on_script_injection_failure(self):
 
         self.assertIsNone(result)
 
-    def test_get_serialized_dom_stitches_cross_origin_iframes(self):
-        """get_serialized_dom stitches cross-origin iframe HTML into srcdoc."""
+    def test_get_serialized_dom_populates_cors_iframes(self):
         driver = Mock()
-        # process_frame calls execute_script twice inside the iframe:
-        #   1. inject percy_dom_script  2. serialize the frame DOM
         driver.execute_script.side_effect = [
             {
                 "html": '<html><body><iframe data-percy-element-id="cid-1"></iframe></body></html>',
-                "resources": [{"url": "https://cdn/main.css", "content": "m"}]
-            },                                # main page serialize (get_serialized_dom)
-            None,                         # inject percy_dom_script into frame
+                "resources": [{"url": "https://cdn/main.css", "content": "m"}]},                            
+            None,
             {
                 "html": "<iframe-html/>",
-                "resources": [{"url": "https://cdn/frame.css", "content": "f"}]
-            },                              # frame DOM serialize
+                "resources": [{"url": "https://cdn/frame.css", "content": "f"}]},                          
         ]
         driver.current_url = "http://main.example.com/"
 
@@ -1022,11 +1017,14 @@ def test_get_serialized_dom_stitches_cross_origin_iframes(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="some_script")
 
-        self.assertNotIn("corsIframes", dom)
-        self.assertIn('data-percy-element-id="cid-1"', dom["html"])
-        self.assertIn('srcdoc="<iframe-html/>"', dom["html"])
-        urls = [resource["url"] for resource in dom["resources"]]
-        self.assertEqual(urls, ["https://cdn/main.css", "https://cdn/frame.css"])
+        self.assertIn("corsIframes", dom)
+        self.assertEqual(len(dom["corsIframes"]), 1)
+        entry = dom["corsIframes"][0]
+        self.assertEqual(entry["iframeData"]["percyElementId"], "cid-1")
+        self.assertEqual(entry["iframeSnapshot"]["html"], "<iframe-html/>")
+        self.assertEqual(entry["frameUrl"], "https://cross.example.com/page")
+        # HTML is left unchanged (no srcdoc injection here — core handles that)
+        self.assertNotIn("srcdoc", dom["html"])
 
     def test_get_serialized_dom_skips_blank_src_frames(self):
         """Frames with no src or src='about:blank' are not processed."""
@@ -1044,7 +1042,7 @@ def test_get_serialized_dom_skips_blank_src_frames(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="some_script")
 
-        self.assertNotIn("srcdoc=", dom["html"])
+        self.assertNotIn("corsIframes", dom)
 
     def test_get_serialized_dom_no_cors_iframes_without_script(self):
         """Without a percy_dom_script, cross-origin iframes are not processed."""
@@ -1062,7 +1060,7 @@ def test_get_serialized_dom_no_cors_iframes_without_script(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script=None)
 
-        self.assertNotIn("srcdoc=", dom["html"])
+        self.assertNotIn("corsIframes", dom)
 
     def test_get_serialized_dom_cookies_always_attached(self):
         """Cookies are always added to the dom_snapshot regardless of iframes."""
@@ -1096,7 +1094,8 @@ def test_get_serialized_dom_same_host_different_scheme_is_cross_origin(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="script")
 
-        self.assertIn('srcdoc="<frame/>"', dom["html"])
+        self.assertIn("corsIframes", dom)
+        self.assertEqual(dom["corsIframes"][0]["iframeSnapshot"]["html"], "<frame/>")
 
     def test_get_serialized_dom_same_host_different_port_is_cross_origin(self):
         """http://example.com:3000 and http://example.com:4000 differ in port → cross-origin."""
@@ -1116,7 +1115,8 @@ def test_get_serialized_dom_same_host_different_port_is_cross_origin(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="script")
 
-        self.assertIn('srcdoc="<frame/>"', dom["html"])
+        self.assertIn("corsIframes", dom)
+        self.assertEqual(dom["corsIframes"][0]["iframeSnapshot"]["html"], "<frame/>")
 
     def test_get_serialized_dom_same_origin_is_not_cross_origin(self):
         """http://main.example.com/page1 and http://main.example.com/page2 share origin."""
@@ -1132,7 +1132,7 @@ def test_get_serialized_dom_same_origin_is_not_cross_origin(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="script")
 
-        self.assertNotIn("srcdoc=", dom["html"])
+        self.assertNotIn("corsIframes", dom)
 
     def test_process_frame_passes_enable_javascript_option(self):
         """process_frame serializes the frame with enableJavaScript=True, mirroring
@@ -1169,31 +1169,30 @@ def test_process_frame_uses_unknown_src_fallback(self):
         self.assertIsNotNone(result)
         self.assertEqual(result["frameUrl"], "unknown-src")
 
-    def test_stitch_cors_iframes_injects_srcdoc_and_merges_resources(self):
-        """stitch_cors_iframes injects iframe HTML into srcdoc and deduplicates resources by URL."""
-        dom_snapshot = {
-            "html": '<html><body><iframe data-percy-element-id="cid-1"></iframe></body></html>',
-            "resources": [{"url": "https://cdn/main.css", "content": "m"}]
-        }
-        processed_frames = [{
-            "iframeData": {"percyElementId": "cid-1"},
-            "iframeSnapshot": {
-                "html": '<html><body><h1>F & "Q"</h1></body></html>',
-                "resources": [
-                    {"url": "https://cdn/main.css", "content": "dup"},
-                    {"url": "https://cdn/frame.css", "content": "f"}
-                ]
-            }
-        }]
+    def test_get_serialized_dom_corsIframes_entry_has_correct_structure(self):
+        dom_html = '<html><body><iframe data-percy-element-id="cid-1"></iframe></body></html>'
+        frame_resource = {"url": "https://cdn/frame.css", "content": "f"}
+        driver = Mock()
+        driver.execute_script.side_effect = [
+            {"html": dom_html, "resources": []},
+            None,
+            {"html": '<html><body><h1>Frame</h1></body></html>', "resources": [frame_resource]},
+        ]
+        driver.current_url = "http://main.example.com/"
+        frame = Mock()
+        frame.get_attribute = lambda attr: (
+            "https://cross.example.com/page" if attr == 'src' else "cid-1"
+        )
+        driver.find_elements.return_value = [frame]
 
-        stitched = local.stitch_cors_iframes(dom_snapshot, processed_frames)
+        dom = local.get_serialized_dom(driver, [], percy_dom_script="some_script")
 
-        self.assertIn(
-            'srcdoc="<html><body><h1>F &amp; &quot;Q&quot;</h1></body></html>"',
-            stitched["html"]
-        )
-        urls = [resource["url"] for resource in stitched["resources"]]
-        self.assertEqual(urls, ["https://cdn/main.css", "https://cdn/frame.css"])
+        self.assertIn("corsIframes", dom)
+        entry = dom["corsIframes"][0]
+        self.assertEqual(entry["frameUrl"], "https://cross.example.com/page")
+        self.assertEqual(entry["iframeData"], {"percyElementId": "cid-1"})
+        self.assertIn("Frame", entry["iframeSnapshot"]["html"])
+        self.assertIn(frame_resource, entry["iframeSnapshot"]["resources"])
 
     def test_get_serialized_dom_multiple_cross_origin_frames(self):
         """All cross-origin frames are collected; same-origin frames are skipped."""
@@ -1225,9 +1224,12 @@ def test_get_serialized_dom_multiple_cross_origin_frames(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="script")
 
-        self.assertIn('data-percy-element-id="pid-1" srcdoc="<frame1/>"', dom["html"])
-        self.assertIn('data-percy-element-id="pid-2" srcdoc="<frame2/>"', dom["html"])
-        self.assertNotIn('data-percy-element-id="pid-same" srcdoc=', dom["html"])
+        self.assertIn("corsIframes", dom)
+        self.assertEqual(len(dom["corsIframes"]), 2)
+        pids = [e["iframeData"]["percyElementId"] for e in dom["corsIframes"]]
+        self.assertIn("pid-1", pids)
+        self.assertIn("pid-2", pids)
+        self.assertNotIn("pid-same", pids)
 
     def test_get_serialized_dom_handles_find_elements_exception(self):
         """If find_elements raises, the error is swallowed, cookies are still attached,
@@ -1240,7 +1242,7 @@ def test_get_serialized_dom_handles_find_elements_exception(self):
         dom = local.get_serialized_dom(driver, [{"name": "k", "value": "v"}],
                                        percy_dom_script="script")
 
-        self.assertNotIn("srcdoc=", dom["html"])
+        self.assertNotIn("corsIframes", dom)
         self.assertEqual(dom["cookies"], [{"name": "k", "value": "v"}])
 
     def test_get_serialized_dom_process_frame_failure_is_skipped(self):
@@ -1270,8 +1272,10 @@ def test_get_serialized_dom_process_frame_failure_is_skipped(self):
 
         dom = local.get_serialized_dom(driver, [], percy_dom_script="script")
 
-        self.assertIn('data-percy-element-id="pid-ok" srcdoc="<ok/>"', dom["html"])
-        self.assertNotIn('data-percy-element-id="pid-fail" srcdoc=', dom["html"])
+        self.assertIn("corsIframes", dom)
+        self.assertEqual(len(dom["corsIframes"]), 1)
+        self.assertEqual(dom["corsIframes"][0]["iframeData"]["percyElementId"], "pid-ok")
+        self.assertEqual(dom["corsIframes"][0]["iframeSnapshot"]["html"], "<ok/>")
 
 
 class TestCreateRegion(unittest.TestCase):