[WebDriver BiDi] extend storage.setCookie tests (#43995)

* Cover `context` partition happy case
* Cover `storageKey` partition happy case
* Refactor cookie assertion

Tracking issue: #43996
---------

Co-authored-by: Maksim Sadym <sadym@google.com>

webdriver/tests/bidi/storage/__init__.py

@@ -0,0 +1,28 @@
+from webdriver.bidi.modules.storage import StorageKeyPartitionDescriptor
+from .. import any_int, recursive_compare
+
+
+async def assert_cookie_is_set(bidi_session, name: str, str_value: str, domain: str, origin: str):
+    """
+    Asserts the cookie is set.
+    """
+    all_cookies = await bidi_session.storage.get_cookies(partition=StorageKeyPartitionDescriptor(
+        source_origin=origin))
+
+    assert 'cookies' in all_cookies
+    cookie = next(c for c in all_cookies['cookies'] if c['name'] == name)
+
+    recursive_compare({
+        'domain': domain,
+        'httpOnly': False,
+        'name': name,
+        'path': '/',
+        'sameSite': 'none',
+        'secure': True,
+        # Varies depending on the cookie name and value.
+        'size': any_int,
+        'value': {
+            'type': 'string',
+            'value': str_value,
+        },
+    }, cookie)

webdriver/tests/bidi/storage/set_cookie/page_protocols.py

@@ -0,0 +1,45 @@
+import pytest
+from webdriver.bidi.modules.network import NetworkStringValue
+from webdriver.bidi.modules.storage import PartialCookie, BrowsingContextPartitionDescriptor
+from .. import assert_cookie_is_set
+
+pytestmark = pytest.mark.asyncio
+
+COOKIE_NAME = 'SOME_COOKIE_NAME'
+COOKIE_VALUE = 'SOME_COOKIE_VALUE'
+
+
+@pytest.mark.parametrize(
+    "protocol",
+    [
+        "http",
+        "https",
+    ]
+)
+async def test_page_protocols(bidi_session, top_context, get_test_page, origin, domain_value, protocol):
+    # Navigate to a page with a required protocol.
+    await bidi_session.browsing_context.navigate(
+        context=top_context["context"], url=get_test_page(protocol=protocol), wait="complete"
+    )
+
+    source_origin = origin(protocol)
+    partition = BrowsingContextPartitionDescriptor(top_context["context"])
+
+    set_cookie_result = await bidi_session.storage.set_cookie(
+        cookie=PartialCookie(
+            name=COOKIE_NAME,
+            value=NetworkStringValue(COOKIE_VALUE),
+            domain=domain_value(),
+            secure=True
+        ),
+        partition=partition)
+
+    assert set_cookie_result == {
+        'partitionKey': {
+            'sourceOrigin': source_origin
+        },
+    }
+
+    # Assert the cookie is actually set.
+    await assert_cookie_is_set(bidi_session, name=COOKIE_NAME, str_value=COOKIE_VALUE,
+                               domain=domain_value(), origin=source_origin)

webdriver/tests/bidi/storage/set_cookie/partition_context.py

@@ -0,0 +1,35 @@
+import pytest
+from webdriver.bidi.modules.network import NetworkStringValue
+from webdriver.bidi.modules.storage import PartialCookie, BrowsingContextPartitionDescriptor
+from .. import assert_cookie_is_set
+
+pytestmark = pytest.mark.asyncio
+
+COOKIE_NAME = 'SOME_COOKIE_NAME'
+COOKIE_VALUE = 'SOME_COOKIE_VALUE'
+
+
+async def test_context_partition(bidi_session, top_context, test_page, origin, domain_value):
+    # Navigate to a secure context.
+    await bidi_session.browsing_context.navigate(context=top_context["context"], url=test_page, wait="complete")
+
+    source_origin = origin()
+    partition = BrowsingContextPartitionDescriptor(top_context["context"])
+
+    set_cookie_result = await bidi_session.storage.set_cookie(
+        cookie=PartialCookie(
+            name=COOKIE_NAME,
+            value=NetworkStringValue(COOKIE_VALUE),
+            domain=domain_value(),
+            secure=True
+        ),
+        partition=partition)
+
+    assert set_cookie_result == {
+        'partitionKey': {
+            'sourceOrigin': source_origin
+        },
+    }
+
+    await assert_cookie_is_set(bidi_session, name=COOKIE_NAME, str_value=COOKIE_VALUE,
+                               domain=domain_value(), origin=source_origin)

webdriver/tests/bidi/storage/set_cookie/partition_storage_key.py

@@ -0,0 +1,37 @@
+import pytest
+from webdriver.bidi.modules.network import NetworkStringValue
+from webdriver.bidi.modules.storage import PartialCookie, StorageKeyPartitionDescriptor
+from .. import assert_cookie_is_set
+
+pytestmark = pytest.mark.asyncio
+
+COOKIE_NAME = 'SOME_COOKIE_NAME'
+COOKIE_VALUE = 'SOME_COOKIE_VALUE'
+
+
+async def test_storage_key_partition_source_origin(bidi_session, top_context, test_page, origin, domain_value):
+    # Navigate to a secure context.
+    await bidi_session.browsing_context.navigate(context=top_context["context"], url=test_page, wait="complete")
+
+    source_origin = origin()
+    partition = StorageKeyPartitionDescriptor(source_origin=source_origin)
+
+    set_cookie_result = await bidi_session.storage.set_cookie(
+        cookie=PartialCookie(
+            name=COOKIE_NAME,
+            value=NetworkStringValue(COOKIE_VALUE),
+            domain=domain_value(),
+            secure=True
+        ),
+        partition=partition)
+
+    assert set_cookie_result == {
+        'partitionKey': {
+            'sourceOrigin': source_origin
+        },
+    }
+
+    await assert_cookie_is_set(bidi_session, name=COOKIE_NAME, str_value=COOKIE_VALUE,
+                               domain=domain_value(), origin=source_origin)
+
+# TODO: test `test_storage_key_partition_user_context`.