use Session class for session module

Author: navin772

py/selenium/webdriver/common/bidi/script.py

@@ -18,8 +18,7 @@
 from dataclasses import dataclass
 from typing import List
 
-from .session import session_subscribe
-from .session import session_unsubscribe
+from .session import Session
 
 
 class Script:
@@ -43,12 +42,14 @@ def remove_console_message_handler(self, id):
 
     def _subscribe_to_log_entries(self):
         if not self.log_entry_subscribed:
-            self.conn.execute(session_subscribe(LogEntryAdded.event_class))
+            session = Session(self.conn)
+            self.conn.execute(session.subscribe(LogEntryAdded.event_class))
             self.log_entry_subscribed = True
 
     def _unsubscribe_from_log_entries(self):
         if self.log_entry_subscribed and LogEntryAdded.event_class not in self.conn.callbacks:
-            self.conn.execute(session_unsubscribe(LogEntryAdded.event_class))
+            session = Session(self.conn)
+            self.conn.execute(session.unsubscribe(LogEntryAdded.event_class))
             self.log_entry_subscribed = False
 
     def _handle_log_entry(self, type, handler):

py/selenium/webdriver/common/bidi/session.py

@@ -15,50 +15,43 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from selenium.webdriver.common.bidi.common import command_builder
 
-def session_subscribe(*events, browsing_contexts=None):
-    cmd_dict = {
-        "method": "session.subscribe",
-        "params": {
-            "events": events,
-        },
-    }
-    if browsing_contexts is None:
-        browsing_contexts = []
-    if browsing_contexts:
-        cmd_dict["params"]["browsingContexts"] = browsing_contexts
-    _ = yield cmd_dict
-    return None
 
+class Session:
+
+    def __init__(self, conn):
+        self.conn = conn
 
-def session_unsubscribe(*events, browsing_contexts=None):
-    cmd_dict = {
-        "method": "session.unsubscribe",
-        "params": {
+    def subscribe(self, *events, browsing_contexts=None):
+        params = {
             "events": events,
-        },
-    }
-    if browsing_contexts is None:
-        browsing_contexts = []
-    if browsing_contexts:
-        cmd_dict["params"]["browsingContexts"] = browsing_contexts
-    _ = yield cmd_dict
-    return None
+        }
+        if browsing_contexts is None:
+            browsing_contexts = []
+        if browsing_contexts:
+            params["browsingContexts"] = browsing_contexts
+        return command_builder("session.subscribe", params)
 
+    def unsubscribe(self, *events, browsing_contexts=None):
+        params = {
+            "events": events,
+        }
+        if browsing_contexts is None:
+            browsing_contexts = []
+        if browsing_contexts:
+            params["browsingContexts"] = browsing_contexts
+        return command_builder("session.unsubscribe", params)
 
-def session_status():
-    """
-    The session.status command returns information about the remote end's readiness
-    to create new sessions and may include implementation-specific metadata.
+    def status(self):
+        """
+        The session.status command returns information about the remote end's readiness
+        to create new sessions and may include implementation-specific metadata.
 
-    Returns
-    -------
-    dict
-        Dictionary containing the ready state (bool), message (str) and metadata
-    """
-    cmd_dict = {
-        "method": "session.status",
-        "params": {},
-    }
-    result = yield cmd_dict
-    return result
+        Returns
+        -------
+        dict
+            Dictionary containing the ready state (bool), message (str) and metadata
+        """
+        cmd = command_builder("session.status", {})
+        return self.conn.execute(cmd)

py/selenium/webdriver/remote/webdriver.py

@@ -44,6 +44,7 @@
 from selenium.webdriver.common.bidi.browser import Browser
 from selenium.webdriver.common.bidi.network import Network
 from selenium.webdriver.common.bidi.script import Script
+from selenium.webdriver.common.bidi.session import Session
 from selenium.webdriver.common.by import By
 from selenium.webdriver.common.options import ArgOptions
 from selenium.webdriver.common.options import BaseOptions
@@ -256,6 +257,7 @@ def __init__(
         self._script = None
         self._network = None
         self._browser = None
+        self._session = None
 
     def __repr__(self):
         return f'<{type(self).__module__}.{type(self).__name__} (session="{self.session_id}")>'
@@ -1294,29 +1296,23 @@ def browser(self):
 
         return self._browser
 
-    def get_bidi_session_status(self):
-        """
-        Get the session status using WebDriver BiDi.
-        Returns information about whether a remote end is in a state
-        in which it can create new sessions.
-
-        Returns:
-        -------
-        dict
-            Dictionary containing the ready state (bool), message (str) and metadata
+    @property
+    def session(self):
+        """Returns the BiDi session object for the current WebDriver session.
 
         Example:
         --------
-        >>> status = driver.get_bidi_session_status()
-        >>> print(status["ready"])
-        >>> print(status["message"])
+        >>> driver.session.subscribe()
+        >>> driver.session.unsubscribe()
+        >>> session = driver.session.status()
         """
         if not self._websocket_connection:
             self._start_bidi()
 
-        from selenium.webdriver.common.bidi.session import session_status
+        if self._session is None:
+            self._session = Session(self._websocket_connection)
 
-        return self._websocket_connection.execute(session_status())
+        return self._session
 
     def _get_cdp_details(self):
         import json

py/test/selenium/webdriver/common/bidi_session_tests.py

@@ -22,8 +22,7 @@
 
 @pytest.mark.xfail_safari
 def test_session_status(driver):
-    result = driver.get_bidi_session_status()
-
+    result = driver.session.status()
     assert result is not None
     assert "ready" in result
     assert "message" in result
@@ -34,7 +33,7 @@ def test_session_status(driver):
 @pytest.mark.xfail_safari
 def test_session_status_not_closed_with_one_window(driver):
     # initial session status
-    initial_status = driver.get_bidi_session_status()
+    initial_status = driver.session.status()
     assert initial_status is not None
 
     # Open new window and tab
@@ -45,7 +44,7 @@ def test_session_status_not_closed_with_one_window(driver):
     driver.close()
 
     # Session should still be active
-    status_after_closing = driver.get_bidi_session_status()
+    status_after_closing = driver.session.status()
     assert status_after_closing is not None
     assert "ready" in status_after_closing
     assert "message" in status_after_closing