feat: addPreloadScript respects new contexts

GoogleChromeLabs · Oct 25, 2023 · ea0b382 · ea0b382
1 parent c3af4d0
commit ea0b382
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 2 deletions.
diff --git a/src/bidiMapper/domains/context/CdpTarget.ts b/src/bidiMapper/domains/context/CdpTarget.ts
@@ -200,7 +200,9 @@ export class CdpTarget {
 
   /** Loads all top-level preload scripts. */
   async #initAndEvaluatePreloadScripts() {
-    for (const script of this.#preloadScriptStorage.find()) {
+    for (const script of this.#preloadScriptStorage.find({
+      global: true,
+    })) {
       await script.initInTarget(this, true);
     }
   }

diff --git a/src/bidiMapper/domains/script/PreloadScriptStorage.ts b/src/bidiMapper/domains/script/PreloadScriptStorage.ts
@@ -20,7 +20,10 @@ import type {PreloadScript} from './PreloadScript.js';
 
 /** PreloadScripts can be filtered by BiDi ID or target ID. */
 export type PreloadScriptFilter = Partial<
-  Pick<PreloadScript, 'id'> & Pick<CdpTarget, 'targetId'>
+  Pick<PreloadScript, 'id'> &
+    Pick<CdpTarget, 'targetId'> & {
+      global: boolean;
+    }
 >;
 
 /**
@@ -46,6 +49,15 @@ export class PreloadScriptStorage {
       ) {
         return false;
       }
+      if (
+        filter.global !== undefined &&
+        // Global scripts have no contexts
+        ((filter.global && script.contexts !== undefined) ||
+          // Non global scripts always have contexts
+          (!filter.global && script.contexts === undefined))
+      ) {
+        return false;
+      }
       return true;
     });
   }

diff --git a/tests/script/test_add_preload_script.py b/tests/script/test_add_preload_script.py
@@ -895,3 +895,64 @@ async def test_preloadScript_add_respectContextsForOldContexts(
             }
         })
     assert result["result"] == {"type": "undefined"}
+
+
+@pytest.mark.asyncio
+async def test_preloadScript_add_respectContextsForNewContexts(
+        websocket, context_id, html):
+
+    # Add the PreloadScript only to a specific context
+    await execute_command(
+        websocket, {
+            "method": "script.addPreloadScript",
+            "params": {
+                "functionDeclaration": """
+                    () => {
+                        window.FOO = "BAR"
+                    }""",
+                "contexts": [context_id]
+            }
+        })
+
+    # Create a new context after adding PreloadScript
+    result = await execute_command(websocket, {
+        "method": "browsingContext.create",
+        "params": {
+            "type": "tab"
+        }
+    })
+    new_context_id = result['context']
+
+    # Navigate both contexts to trigger PreloadScripts
+    await goto_url(websocket, context_id, html())
+    await goto_url(websocket, new_context_id, html())
+
+    # Expect context with context_id to be affected by PreloadScript
+    result = await execute_command(
+        websocket, {
+            "method": "script.evaluate",
+            "params": {
+                "expression": "window.FOO",
+                "target": {
+                    "context": context_id
+                },
+                "awaitPromise": True,
+                "resultOwnership": "root"
+            }
+        })
+    assert result["result"] == {"type": "string", "value": 'BAR'}
+
+    # Expect context with new_context_id to not be affected by PreloadScript
+    result = await execute_command(
+        websocket, {
+            "method": "script.evaluate",
+            "params": {
+                "expression": "window.FOO",
+                "target": {
+                    "context": new_context_id
+                },
+                "awaitPromise": True,
+                "resultOwnership": "root"
+            }
+        })
+    assert result["result"] == {"type": "undefined"}