[py][bidi]: add set_screen_orientation_override command in Emulation

py/selenium/webdriver/common/bidi/emulation.py

@@ -15,11 +15,66 @@
 # specific language governing permissions and limitations
 # under the License.
 
+from enum import Enum
 from typing import Any, Optional, Union
 
 from selenium.webdriver.common.bidi.common import command_builder
 
 
+class ScreenOrientationNatural(Enum):
+    """Natural screen orientation."""
+
+    PORTRAIT = "portrait"
+    LANDSCAPE = "landscape"
+
+
+class ScreenOrientationType(Enum):
+    """Screen orientation type."""
+
+    PORTRAIT_PRIMARY = "portrait-primary"
+    PORTRAIT_SECONDARY = "portrait-secondary"
+    LANDSCAPE_PRIMARY = "landscape-primary"
+    LANDSCAPE_SECONDARY = "landscape-secondary"
+
+
+def _convert_to_enum(value, enum_class):
+    if isinstance(value, enum_class):
+        return value
+    try:
+        return enum_class(value)
+    except ValueError:
+        raise ValueError(f"Invalid orientation: {value}")
+
+
+class ScreenOrientation:
+    """Represents screen orientation configuration."""
+
+    def __init__(
+        self,
+        natural: Union[ScreenOrientationNatural, str],
+        type: Union[ScreenOrientationType, str],
+    ):
+        """Initialize ScreenOrientation.
+
+        Args:
+            natural: Natural screen orientation ("portrait" or "landscape").
+            type: Screen orientation type ("portrait-primary", "portrait-secondary",
+                "landscape-primary", or "landscape-secondary").
+
+        Raises:
+            ValueError: If natural or type values are invalid.
+        """
+        # handle string values
+        self.natural = _convert_to_enum(natural, ScreenOrientationNatural)
+        self.type = _convert_to_enum(type, ScreenOrientationType)
+
+    def to_dict(self) -> dict[str, str]:
+        return {
+            "natural": self.natural.value,
+            "type": self.type.value,
+        }
+
+
 class GeolocationCoordinates:
     """Represents geolocation coordinates."""
 
@@ -310,3 +365,37 @@ def set_scripting_enabled(
             params["userContexts"] = user_contexts
 
         self.conn.execute(command_builder("emulation.setScriptingEnabled", params))
+
+    def set_screen_orientation_override(
+        self,
+        screen_orientation: Optional[ScreenOrientation] = None,
+        contexts: Optional[list[str]] = None,
+        user_contexts: Optional[list[str]] = None,
+    ) -> None:
+        """Set screen orientation override for the given contexts or user contexts.
+
+        Args:
+            screen_orientation: ScreenOrientation object to emulate, or None to clear the override.
+            contexts: List of browsing context IDs to apply the override to.
+            user_contexts: List of user context IDs to apply the override to.
+
+        Raises:
+            ValueError: If both contexts and user_contexts are provided, or if neither
+                contexts nor user_contexts are provided.
+        """
+        if contexts is not None and user_contexts is not None:
+            raise ValueError("Cannot specify both contexts and userContexts")
+
+        if contexts is None and user_contexts is None:
+            raise ValueError("Must specify either contexts or userContexts")
+
+        params: dict[str, Any] = {
+            "screenOrientation": screen_orientation.to_dict() if screen_orientation is not None else None
+        }
+
+        if contexts is not None:
+            params["contexts"] = contexts
+        elif user_contexts is not None:
+            params["userContexts"] = user_contexts
+
+        self.conn.execute(command_builder("emulation.setScreenOrientationOverride", params))

py/test/selenium/webdriver/common/bidi_emulation_tests.py

@@ -16,7 +16,14 @@
 # under the License.
 import pytest
 
-from selenium.webdriver.common.bidi.emulation import Emulation, GeolocationCoordinates, GeolocationPositionError
+from selenium.webdriver.common.bidi.emulation import (
+    Emulation,
+    GeolocationCoordinates,
+    GeolocationPositionError,
+    ScreenOrientation,
+    ScreenOrientationNatural,
+    ScreenOrientationType,
+)
 from selenium.webdriver.common.bidi.permissions import PermissionState
 from selenium.webdriver.common.window import WindowTypes
 
@@ -73,6 +80,24 @@ def get_browser_locale(driver):
     return result.result["value"]
 
 
+def get_screen_orientation(driver, context_id):
+    result = driver.script._evaluate(
+        "screen.orientation.type",
+        {"context": context_id},
+        await_promise=False,
+    )
+    orientation_type = result.result["value"]
+
+    result = driver.script._evaluate(
+        "screen.orientation.angle",
+        {"context": context_id},
+        await_promise=False,
+    )
+    orientation_angle = result.result["value"]
+
+    return {"type": orientation_type, "angle": orientation_angle}
+
+
 def test_emulation_initialized(driver):
     assert driver.emulation is not None
     assert isinstance(driver.emulation, Emulation)
@@ -419,3 +444,86 @@ def test_set_scripting_enabled_with_user_contexts(driver, pages):
             driver.browsing_context.close(context_id)
     finally:
         driver.browser.remove_user_context(user_context)
+
+
+def test_set_screen_orientation_override_with_contexts(driver, pages):
+    context_id = driver.current_window_handle
+    initial_orientation = get_screen_orientation(driver, context_id)
+
+    # Set landscape-primary orientation
+    orientation = ScreenOrientation(
+        natural=ScreenOrientationNatural.LANDSCAPE,
+        type=ScreenOrientationType.LANDSCAPE_PRIMARY,
+    )
+    driver.emulation.set_screen_orientation_override(screen_orientation=orientation, contexts=[context_id])
+
+    url = pages.url("formPage.html")
+    driver.browsing_context.navigate(context_id, url, wait="complete")
+
+    # Verify the orientation was set
+    current_orientation = get_screen_orientation(driver, context_id)
+    assert current_orientation["type"] == "landscape-primary", f"Expected landscape-primary, got {current_orientation}"
+    assert current_orientation["angle"] == 0, f"Expected angle 0, got {current_orientation['angle']}"
+
+    # Set portrait-secondary orientation
+    orientation = ScreenOrientation(
+        natural=ScreenOrientationNatural.PORTRAIT,
+        type=ScreenOrientationType.PORTRAIT_SECONDARY,
+    )
+    driver.emulation.set_screen_orientation_override(screen_orientation=orientation, contexts=[context_id])
+
+    # Verify the orientation was changed
+    current_orientation = get_screen_orientation(driver, context_id)
+    assert current_orientation["type"] == "portrait-secondary", (
+        f"Expected portrait-secondary, got {current_orientation}"
+    )
+    assert current_orientation["angle"] == 180, f"Expected angle 180, got {current_orientation['angle']}"
+
+    driver.emulation.set_screen_orientation_override(screen_orientation=None, contexts=[context_id])
+
+    # Verify orientation was cleared
+    assert get_screen_orientation(driver, context_id) == initial_orientation
+
+
+@pytest.mark.parametrize(
+    "natural,orientation_type,expected_angle",
+    [
+        # Portrait natural orientations
+        ("portrait", "portrait-primary", 0),
+        ("portrait", "portrait-secondary", 180),
+        ("portrait", "landscape-primary", 90),
+        ("portrait", "landscape-secondary", 270),
+        # Landscape natural orientations
+        ("landscape", "portrait-primary", 90),
+        ("landscape", "portrait-secondary", 270),
+        ("landscape", "landscape-primary", 0),
+        ("landscape", "landscape-secondary", 180),
+    ],
+)
+def test_set_screen_orientation_override_with_user_contexts(driver, pages, natural, orientation_type, expected_angle):
+    user_context = driver.browser.create_user_context()
+    try:
+        context_id = driver.browsing_context.create(type=WindowTypes.TAB, user_context=user_context)
+        try:
+            driver.switch_to.window(context_id)
+
+            # Set the specified orientation
+            orientation = ScreenOrientation(natural=natural, type=orientation_type)
+            driver.emulation.set_screen_orientation_override(
+                screen_orientation=orientation, user_contexts=[user_context]
+            )
+
+            url = pages.url("formPage.html")
+            driver.browsing_context.navigate(context_id, url, wait="complete")
+
+            # Verify the orientation was set
+            current_orientation = get_screen_orientation(driver, context_id)
+
+            assert current_orientation["type"] == orientation_type
+            assert current_orientation["angle"] == expected_angle
+
+            driver.emulation.set_screen_orientation_override(screen_orientation=None, user_contexts=[user_context])
+        finally:
+            driver.browsing_context.close(context_id)
+    finally:
+        driver.browser.remove_user_context(user_context)