1 files changed, 150 insertions, 0 deletions
diff --git a/testing/web-platform/mozilla/tests/webdriver/bidi/storage/set_cookie/partition.py b/testing/web-platform/mozilla/tests/webdriver/bidi/storage/set_cookie/partition.py
new file mode 100644
index 0000000000..f8e2823dbc
--- /dev/null
+++ b/testing/web-platform/mozilla/tests/webdriver/bidi/storage/set_cookie/partition.py
@@ -0,0 +1,150 @@
+import pytest
+from tests.bidi import recursive_compare
+from tests.support.helpers import get_origin_from_url
+from webdriver.bidi.modules.network import NetworkStringValue
+from webdriver.bidi.modules.storage import (
+    BrowsingContextPartitionDescriptor,
+    PartialCookie,
+)
+
+pytestmark = pytest.mark.asyncio
+
+
+async def test_partition_context(
+    bidi_session,
+    set_cookie,
+    top_context,
+    new_tab,
+    test_page,
+    test_page_cross_origin,
+    domain_value,
+):
+    await bidi_session.browsing_context.navigate(
+        context=new_tab["context"], url=test_page, wait="complete"
+    )
+    source_origin_1 = get_origin_from_url(test_page)
+    await bidi_session.browsing_context.navigate(
+        context=top_context["context"], url=test_page_cross_origin, wait="complete"
+    )
+    source_origin_2 = get_origin_from_url(test_page_cross_origin)
+
+    cookie_name = "foo"
+    cookie_value = "bar"
+    cookie_domain = domain_value()
+    new_tab_partition = BrowsingContextPartitionDescriptor(new_tab["context"])
+
+    set_cookie_result = await set_cookie(
+        cookie=PartialCookie(
+            domain=cookie_domain,
+            name=cookie_name,
+            value=NetworkStringValue(cookie_value),
+        ),
+        partition=new_tab_partition,
+    )
+
+    assert set_cookie_result == {"partitionKey": {"sourceOrigin": source_origin_1}}
+
+    # Check that added cookies are present on the right context.
+    cookies = await bidi_session.storage.get_cookies(partition=new_tab_partition)
+
+    recursive_compare(
+        {
+            "cookies": [
+                {
+                    "domain": cookie_domain,
+                    "httpOnly": False,
+                    "name": cookie_name,
+                    "path": "/",
+                    "sameSite": "none",
+                    "secure": False,
+                    "size": 6,
+                    "value": {"type": "string", "value": cookie_value},
+                }
+            ],
+            "partitionKey": {"sourceOrigin": source_origin_1},
+        },
+        cookies,
+    )
+
+    # Check that added cookies are not present on the other context.
+    cookies = await bidi_session.storage.get_cookies(
+        partition=BrowsingContextPartitionDescriptor(top_context["context"])
+    )
+
+    recursive_compare(
+        {"cookies": [], "partitionKey": {"sourceOrigin": source_origin_2}}, cookies
+    )
+
+
+@pytest.mark.parametrize("domain", ["", "alt"], ids=["same_origin", "cross_origin"])
+async def test_partition_context_iframe(
+    bidi_session, new_tab, inline, domain_value, domain, set_cookie
+):
+    iframe_url = inline("<div id='in-iframe'>foo</div>", domain=domain)
+    source_origin_for_iframe = get_origin_from_url(iframe_url)
+    page_url = inline(f"<iframe src='{iframe_url}'></iframe>")
+    source_origin_for_page = get_origin_from_url(page_url)
+    await bidi_session.browsing_context.navigate(
+        context=new_tab["context"], url=page_url, wait="complete"
+    )
+
+    contexts = await bidi_session.browsing_context.get_tree(root=new_tab["context"])
+    iframe_context = contexts[0]["children"][0]
+    iframe_partition = BrowsingContextPartitionDescriptor(iframe_context["context"])
+
+    cookie_name = "foo"
+    cookie_value = "bar"
+    await set_cookie(
+        cookie=PartialCookie(
+            domain=domain_value(domain),
+            name=cookie_name,
+            value=NetworkStringValue(cookie_value),
+        ),
+        partition=iframe_partition,
+    )
+
+    # Check that added cookies are present on the right context
+    cookies = await bidi_session.storage.get_cookies(partition=iframe_partition)
+
+    expected_cookies = [
+        {
+            "domain": domain_value(domain=domain),
+            "httpOnly": False,
+            "name": cookie_name,
+            "path": "/",
+            "sameSite": "none",
+            "secure": False,
+            "size": 6,
+            "value": {"type": "string", "value": cookie_value},
+        }
+    ]
+    recursive_compare(
+        {
+            "cookies": expected_cookies,
+            "partitionKey": {"sourceOrigin": source_origin_for_iframe},
+        },
+        cookies,
+    )
+
+    cookies = await bidi_session.storage.get_cookies(
+        partition=BrowsingContextPartitionDescriptor(new_tab["context"])
+    )
+    # When the iframe is on the different domain we can verify that top context has no iframe cookie.
+    if domain == "alt":
+        recursive_compare(
+            {
+                "cookies": [],
+                "partitionKey": {"sourceOrigin": source_origin_for_page},
+            },
+            cookies,
+        )
+    else:
+        # When the iframe is on the same domain, since the browsing context partition is defined by user context and origin,
+        # which will be the same for the page, we get the same cookies as for the iframe
+        recursive_compare(
+            {
+                "cookies": expected_cookies,
+                "partitionKey": {"sourceOrigin": source_origin_for_page},
+            },
+            cookies,
+        )