fix(webdriver): support W3C appium: prefix and custom connection for directConnect

README.md

@@ -291,12 +291,12 @@ any | `appium.options.common.base.AppiumOptions`
 
 ## Direct Connect URLs
 
-If your Selenium/Appium server decorates the new session capabilities response with the following keys:
+If your Selenium/Appium server decorates the new session capabilities response with the following keys (with or without the `appium:` prefix):
 
-- `directConnectProtocol`
-- `directConnectHost`
-- `directConnectPort`
-- `directConnectPath`
+- `directConnectProtocol` (`appium:directConnectProtocol`)
+- `directConnectHost` (`appium:directConnectHost`)
+- `directConnectPort` (`appium:directConnectPort`)
+- `directConnectPath` (`appium:directConnectPath`)
 
 Then python client will switch its endpoint to the one specified by the values of those keys.
 

appium/webdriver/webdriver.py

@@ -28,7 +28,7 @@
 from selenium.webdriver.remote.webdriver import WebDriver as Remote
 
 from appium.common.logger import logger
-from appium.options.common.base import AppiumOptions
+from appium.options.common.base import APPIUM_PREFIX, AppiumOptions
 
 from .appium_connection import AppiumConnection
 from .client_config import AppiumClientConfig
@@ -213,6 +213,14 @@ def _get_remote_connection_and_client_config(
     return (AppiumConnection(client_config=new_client_config), new_client_config)
 
 
+def _get_direct_connect_value(caps: dict[str, Any], key: str) -> Any:
+    for candidate in (f'{APPIUM_PREFIX}{key}', key):
+        value = caps.get(candidate)
+        if value is not None and value != '':
+            return value
+    return None
+
+
 class WebDriver(
     Remote,
     ActionHelpers,
@@ -273,8 +281,9 @@ def __init__(
 
         self.error_handler = MobileErrorHandler()
 
-        if client_config and client_config.direct_connection:
-            self._update_command_executor(keep_alive=client_config.keep_alive)
+        effective_client_config = client_config or getattr(self.command_executor, 'client_config', None)
+        if isinstance(effective_client_config, AppiumClientConfig) and effective_client_config.direct_connection:
+            self._update_command_executor(keep_alive=effective_client_config.keep_alive)
 
         self._absent_extensions: set[str] = set()
 
@@ -316,15 +325,16 @@ def _update_command_executor(self, keep_alive: bool) -> None:
         if not self.caps:
             raise ValueError('Driver capabilities must be defined')
         keys = (direct_protocol, direct_host, direct_port, direct_path)
-        if not set(keys).issubset(self.caps):
-            details = ' '.join(f"{key}: '{self.caps.get(key, '')}'" for key in keys)
+        values = {key: _get_direct_connect_value(self.caps, key) for key in keys}
+        if any(val is None for val in values.values()):
+            details = ' '.join(f"{key}: '{values[key] if values[key] is not None else ''}'" for key in keys)
             logger.debug(f'Direct connect capabilities from server were:\n{details} ')
             return
 
-        protocol = self.caps[direct_protocol]
-        hostname = self.caps[direct_host]
-        port = self.caps[direct_port]
-        path = self.caps[direct_path]
+        protocol = values[direct_protocol]
+        hostname = values[direct_host]
+        port = values[direct_port]
+        path = values[direct_path]
         executor = f'{protocol}://{hostname}:{port}{path}'
 
         logger.debug('Updated request endpoint to %s', executor)
@@ -338,7 +348,7 @@ def _update_command_executor(self, keep_alive: bool) -> None:
         client_config.remote_server_addr = executor
         client_config.keep_alive = keep_alive
         if isinstance(self.command_executor, AppiumConnection):
-            self.command_executor = AppiumConnection(client_config=client_config)
+            self.command_executor = type(self.command_executor)(client_config=client_config)
         else:
             self.command_executor = RemoteConnection(client_config=client_config)
         self._add_commands()

test/unit/webdriver/webdriver_test.py

@@ -234,6 +234,158 @@ def test_create_session_register_uridirect_no_direct_connect_path(self):
         assert driver.contexts == ['NATIVE_APP', 'CHROMIUM']
         assert isinstance(driver.command_executor, AppiumConnection)
 
+    @httpretty.activate
+    def test_create_session_register_uridirect_with_appium_prefix(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body=json.dumps(
+                {
+                    'sessionId': 'session-id',
+                    'capabilities': {
+                        'deviceName': 'Android Emulator',
+                        'appium:directConnectProtocol': 'https',
+                        'appium:directConnectHost': 'localhost2',
+                        'appium:directConnectPort': 4800,
+                        'appium:directConnectPath': '/special/path/wd/hub',
+                    },
+                }
+            ),
+        )
+
+        httpretty.register_uri(
+            httpretty.GET,
+            'https://localhost2:4800/special/path/wd/hub/session/session-id/contexts',
+            body=json.dumps({'value': ['NATIVE_APP', 'CHROMIUM']}),
+        )
+
+        desired_caps = {
+            'platformName': 'Android',
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+            'automationName': 'UIAutomator2',
+        }
+        client_config = AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True)
+        driver = webdriver.Remote(
+            SERVER_URL_BASE,
+            options=UiAutomator2Options().load_capabilities(desired_caps),
+            client_config=client_config,
+        )
+
+        assert driver.command_executor._client_config.remote_server_addr == 'https://localhost2:4800/special/path/wd/hub'
+        assert driver.contexts == ['NATIVE_APP', 'CHROMIUM']
+        assert isinstance(driver.command_executor, AppiumConnection)
+
+    @httpretty.activate
+    def test_create_session_register_uridirect_prefers_appium_prefix(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body=json.dumps(
+                {
+                    'sessionId': 'session-id',
+                    'capabilities': {
+                        'deviceName': 'Android Emulator',
+                        'directConnectProtocol': 'http',
+                        'directConnectHost': 'legacy-host',
+                        'directConnectPort': 4723,
+                        'directConnectPath': '/wd/hub',
+                        'appium:directConnectProtocol': 'https',
+                        'appium:directConnectHost': 'w3c-host',
+                        'appium:directConnectPort': 4800,
+                        'appium:directConnectPath': '/w3c/path',
+                    },
+                }
+            ),
+        )
+
+        desired_caps = {
+            'platformName': 'Android',
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+            'automationName': 'UIAutomator2',
+        }
+        client_config = AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True)
+        driver = webdriver.Remote(
+            SERVER_URL_BASE,
+            options=UiAutomator2Options().load_capabilities(desired_caps),
+            client_config=client_config,
+        )
+
+        assert driver.command_executor._client_config.remote_server_addr == 'https://w3c-host:4800/w3c/path'
+
+    @httpretty.activate
+    def test_create_session_register_uridirect_ignores_empty_values(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body=json.dumps(
+                {
+                    'sessionId': 'session-id',
+                    'capabilities': {
+                        'deviceName': 'Android Emulator',
+                        'appium:directConnectProtocol': 'http',
+                        'appium:directConnectHost': '',
+                        'appium:directConnectPort': 4800,
+                        'appium:directConnectPath': None,
+                    },
+                }
+            ),
+        )
+
+        desired_caps = {
+            'platformName': 'Android',
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+            'automationName': 'UIAutomator2',
+        }
+        client_config = AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True)
+        driver = webdriver.Remote(
+            SERVER_URL_BASE,
+            options=UiAutomator2Options().load_capabilities(desired_caps),
+            client_config=client_config,
+        )
+
+        assert driver.command_executor._client_config.remote_server_addr == SERVER_URL_BASE
+
+    @httpretty.activate
+    def test_create_session_register_uridirect_with_custom_connection(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body=json.dumps(
+                {
+                    'sessionId': 'session-id',
+                    'capabilities': {
+                        'deviceName': 'Android Emulator',
+                        'appium:directConnectProtocol': 'http',
+                        'appium:directConnectHost': 'localhost2',
+                        'appium:directConnectPort': 4800,
+                        'appium:directConnectPath': '/special/path/wd/hub',
+                    },
+                }
+            ),
+        )
+
+        class CustomAppiumConnection(AppiumConnection):
+            pass
+
+        desired_caps = {
+            'platformName': 'Android',
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+            'automationName': 'UIAutomator2',
+        }
+        client_config = AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True)
+        custom_connection = CustomAppiumConnection(client_config=client_config)
+        driver = webdriver.Remote(
+            custom_connection,
+            options=UiAutomator2Options().load_capabilities(desired_caps),
+        )
+
+        assert isinstance(driver.command_executor, CustomAppiumConnection)
+        assert driver.command_executor._client_config.remote_server_addr == 'http://localhost2:4800/special/path/wd/hub'
+
     @httpretty.activate
     def test_create_session_remote_server_addr_treatment_with_appiumclientconfig(self):
         # remote server add in AppiumRemoteCong will be prior than the string of 'command_executor'