[wdspec] add user context subscription tests (#50434)

* [wdspec] add user context subscription tests

* bring back the loop

webdriver/tests/bidi/browsing_context/context_created/context_created.py

@@ -326,3 +326,23 @@ async def test_existing_context(bidi_session, wait_for_event, wait_for_future_sa
         parent=None,
         user_context="default"
     )
+
+
+@pytest.mark.parametrize("type_hint", ["tab", "window"])
+async def test_existing_context_via_user_context(bidi_session, create_user_context, wait_for_event, wait_for_future_safe, subscribe_events, type_hint):
+    user_context = await create_user_context()
+    # See https://w3c.github.io/webdriver-bidi/#ref-for-remote-end-subscribe-steps%E2%91%A1.
+    top_level_context = await bidi_session.browsing_context.create(type_hint=type_hint, user_context=user_context)
+
+    on_entry = wait_for_event(CONTEXT_CREATED_EVENT)
+    await subscribe_events([CONTEXT_CREATED_EVENT], user_contexts=[user_context])
+    context_info = await wait_for_future_safe(on_entry)
+
+    assert_browsing_context(
+        context_info,
+        top_level_context["context"],
+        children=None,
+        url="about:blank",
+        parent=None,
+        user_context="default"
+    )

webdriver/tests/bidi/session/subscribe/user_contexts.py

@@ -0,0 +1,84 @@
+import pytest
+
+from tests.support.sync import AsyncPoll
+from ... import create_console_api_message, recursive_compare
+
+
+@pytest.mark.asyncio
+async def test_subscribe_one_user_context(bidi_session, subscribe_events, create_user_context):
+    user_context = await create_user_context()
+
+    default_context = await bidi_session.browsing_context.create(
+        type_hint="tab",
+        user_context="default"
+    )
+
+    other_context = await bidi_session.browsing_context.create(
+        type_hint="tab",
+        user_context=user_context
+    )
+
+    await subscribe_events(events=["log.entryAdded"], user_contexts=[user_context])
+
+    # Track all received log.entryAdded events in the events array
+    events = []
+
+    async def on_event(method, data):
+        events.append(data)
+
+    remove_listener = bidi_session.add_event_listener("log.entryAdded", on_event)
+
+    await create_console_api_message(bidi_session, default_context, "text1")
+    await create_console_api_message(bidi_session, other_context, "text2")
+
+    wait = AsyncPoll(
+        bidi_session, message="Didn't receive expected log events"
+    )
+    await wait.until(lambda _: len(events) >= 1)
+
+    assert len(events) == 1
+    recursive_compare(
+        {
+            "text": "text2",
+        },
+        events[0],
+    )
+
+    remove_listener()
+
+
+@pytest.mark.asyncio
+async def test_subscribe_multiple_user_contexts(bidi_session, subscribe_events, create_user_context):
+    user_context = await create_user_context()
+
+    default_context = await bidi_session.browsing_context.create(
+        type_hint="tab",
+        user_context="default"
+    )
+
+    other_context = await bidi_session.browsing_context.create(
+        type_hint="tab",
+        user_context=user_context
+    )
+
+    await subscribe_events(events=["log.entryAdded"], user_contexts=[user_context, "default"])
+
+    # Track all received log.entryAdded events in the events array
+    events = []
+
+    async def on_event(method, data):
+        events.append(data)
+
+    remove_listener = bidi_session.add_event_listener("log.entryAdded", on_event)
+
+    await create_console_api_message(bidi_session, default_context, "text1")
+    await create_console_api_message(bidi_session, other_context, "text2")
+
+    wait = AsyncPoll(
+        bidi_session, message="Didn't receive expected log events"
+    )
+    await wait.until(lambda _: len(events) >= 2)
+
+    assert len(events) == 2
+
+    remove_listener()

webdriver/tests/support/fixtures_bidi.py

@@ -65,18 +65,17 @@ async def execute_as_async(sync_func, **kwargs):
 async def subscribe_events(bidi_session):
     subscriptions = []
 
-    async def subscribe_events(events, contexts=None):
-        result = await bidi_session.session.subscribe(events=events, contexts=contexts)
-        subscriptions.append((events, contexts))
+    async def subscribe_events(events, contexts=None, user_contexts=None):
+        result = await bidi_session.session.subscribe(events=events, contexts=contexts, user_contexts=user_contexts)
+        subscriptions.append(result["subscription"])
         return result
 
     yield subscribe_events
 
-    for events, contexts in reversed(subscriptions):
+    for subscription in reversed(subscriptions):
         try:
-            await bidi_session.session.unsubscribe(events=events,
-                                                   contexts=contexts)
-        except (InvalidArgumentException, NoSuchFrameException):
+            await bidi_session.session.unsubscribe(subscriptions=[subscription])
+        except InvalidArgumentException:
             pass