[WebDriver BiDi] addPreloadScript supports contexts

tools/webdriver/webdriver/bidi/modules/script.py

@@ -96,6 +96,7 @@ def add_preload_script(
         self,
         function_declaration: str,
         arguments: Optional[List[Mapping[str, Any]]] = None,
+        contexts: Optional[List[str]] = None,
         sandbox: Optional[str] = None
     ) -> Mapping[str, Any]:
         params: MutableMapping[str, Any] = {
@@ -104,6 +105,8 @@ def add_preload_script(
 
         if arguments is not None:
             params["arguments"] = arguments
+        if contexts is not None:
+            params["contexts"] = contexts
         if sandbox is not None:
             params["sandbox"] = sandbox
 

webdriver/tests/bidi/script/add_preload_script/add_preload_script.py

@@ -170,3 +170,75 @@ async def test_page_script_can_access_preload_script_properties(
         await_promise=True,
     )
     assert result == {"type": "number", "value": 42}
+
+
+@pytest.mark.asyncio
+async def test_add_preload_script_in_iframe_with_top_context_specified(
+    bidi_session, add_preload_script, new_tab, test_page_same_origin_frame
+):
+    await add_preload_script(function_declaration="() => { window.bar='foo'; }", contexts=[new_tab["context"]])
+
+    await bidi_session.browsing_context.navigate(
+        context=new_tab["context"],
+        url=test_page_same_origin_frame,
+        wait="complete",
+    )
+
+    # Check that preload script applied the changes to the window
+    result = await bidi_session.script.evaluate(
+        expression="window.bar",
+        target=ContextTarget(new_tab["context"]),
+        await_promise=True,
+    )
+    assert result == {"type": "string", "value": "foo"}
+
+    contexts = await bidi_session.browsing_context.get_tree(root=new_tab["context"])
+
+    assert len(contexts[0]["children"]) == 1
+    frame_context = contexts[0]["children"][0]
+
+    # Check that preload script applied the changes to the iframe
+    result = await bidi_session.script.evaluate(
+        expression="window.bar",
+        target=ContextTarget(frame_context["context"]),
+        await_promise=True,
+    )
+    assert result == {"type": "string", "value": "foo"}
+
+
+@pytest.mark.asyncio
+async def test_page_script_context_isolation(
+    bidi_session, add_preload_script, top_context, new_tab, test_page
+):
+    await add_preload_script(
+        function_declaration="() => { window.baz = 42; }",
+        contexts=[top_context['context']]
+    )
+
+    # Navigate both contexts to ensure preload script is triggered
+    await bidi_session.browsing_context.navigate(
+        context=top_context['context'],
+        url=test_page,
+        wait="complete",
+    )
+    await bidi_session.browsing_context.navigate(
+        context=new_tab["context"],
+        url=test_page,
+        wait="complete",
+    )
+
+    # Check that preload script applied the changes to the window
+    result = await bidi_session.script.evaluate(
+        expression="window.baz",
+        target=ContextTarget(top_context["context"]),
+        await_promise=True,
+    )
+    assert result == {"type": "number", "value": 42}
+
+    # Check that preload script did *not* apply the changes to the window
+    result = await bidi_session.script.evaluate(
+        expression="window.baz",
+        target=ContextTarget(new_tab["context"]),
+        await_promise=True,
+    )
+    assert result == {type: "undefined"}

webdriver/tests/bidi/script/add_preload_script/invalid.py

@@ -71,7 +71,8 @@ async def test_params_arguments_channel_ownership_invalid_value(bidi_session):
     with pytest.raises(error.InvalidArgumentException):
         await bidi_session.script.add_preload_script(
             function_declaration="() => {}",
-            arguments=[{"type": "channel", "value": {"ownership": "_UNKNOWN_"}}],
+            arguments=[{"type": "channel", "value": {
+                "ownership": "_UNKNOWN_"}}],
         )
 
 
@@ -186,6 +187,43 @@ async def test_params_arguments_channel_include_shadow_tree_invalid_value(bidi_s
         )
 
 
+@pytest.mark.parametrize("contexts", [False, 42, '_UNKNOWN_', {}])
+async def test_params_contexts_invalid_type(bidi_session, contexts):
+    with pytest.raises(error.InvalidArgumentException):
+        await bidi_session.script.add_preload_script(
+            function_declaration="() => {}",
+            contexts=contexts
+        ),
+
+
+async def test_params_contexts_context_invalid_value(bidi_session):
+    with pytest.raises(error.InvalidArgumentException):
+        await bidi_session.script.add_preload_script(
+            function_declaration="() => {}",
+            contexts=["_UNKNOWN_"]
+        ),
+
+
+async def test_params_contexts_context_non_top_level(bidi_session, new_tab, test_page_same_origin_frame):
+    await bidi_session.browsing_context.navigate(
+        context=new_tab["context"],
+        url=test_page_same_origin_frame,
+        wait="complete",
+    )
+
+    contexts = await bidi_session.browsing_context.get_tree(root=new_tab["context"])
+
+    assert len(contexts) == 1
+    assert len(contexts[0]["children"]) == 1
+    child_info = contexts[0]["children"][0]
+
+    with pytest.raises(error.InvalidArgumentException):
+        await bidi_session.script.add_preload_script(
+            function_declaration="() => {}",
+            contexts=[child_info['context']]
+        ),
+
+
 @pytest.mark.parametrize("sandbox", [False, 42, {}, []])
 async def test_params_sandbox_invalid_type(bidi_session, sandbox):
     with pytest.raises(error.InvalidArgumentException):

webdriver/tests/support/fixtures_bidi.py

@@ -18,11 +18,12 @@
 async def add_preload_script(bidi_session):
     preload_scripts_ids = []
 
-    async def add_preload_script(function_declaration, arguments=None, sandbox=None):
+    async def add_preload_script(function_declaration, arguments=None, contexts=None, sandbox=None):
         script = await bidi_session.script.add_preload_script(
             function_declaration=function_declaration,
             arguments=arguments,
             sandbox=sandbox,
+            contexts=contexts,
         )
         preload_scripts_ids.append(script)
 
@@ -214,8 +215,10 @@ async def assert_pdf_dimensions(pdf, expected_dimensions):
         width, height = png_dimensions(png)
 
         # account for potential rounding errors
-        assert (height - 1) <= cm_to_px(expected_dimensions["height"]) <= (height + 1)
-        assert (width - 1) <= cm_to_px(expected_dimensions["width"]) <= (width + 1)
+        assert (
+            height - 1) <= cm_to_px(expected_dimensions["height"]) <= (height + 1)
+        assert (
+            width - 1) <= cm_to_px(expected_dimensions["width"]) <= (width + 1)
 
     return assert_pdf_dimensions
 
@@ -277,7 +280,8 @@ async def compare_png_bidi(img1, img2):
         )
         await bidi_session.browsing_context.close(context=context["context"])
         assert result["type"] == "object"
-        assert set(item[0] for item in result["value"]) == {"totalPixels", "maxDifference"}
+        assert set(item[0] for item in result["value"]) == {
+            "totalPixels", "maxDifference"}
         for item in result["value"]:
             assert len(item) == 2
             assert item[1]["type"] == "number"