fix(webdriver): preserve direct connect root paths and custom executor state (#1365)

Author: KazuCocoa

appium/webdriver/webdriver.py

@@ -216,7 +216,7 @@ def _get_remote_connection_and_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 != '':
+        if value is not None and (value != '' or key == 'directConnectPath'):
             return value
     return None
 
@@ -281,7 +281,7 @@ def __init__(
 
         self.error_handler = MobileErrorHandler()
 
-        effective_client_config = client_config or getattr(self.command_executor, 'client_config', None)
+        effective_client_config = self.command_executor.client_config
         if isinstance(effective_client_config, AppiumClientConfig) and effective_client_config.direct_connection:
             self._update_command_executor(keep_alive=effective_client_config.keep_alive)
 
@@ -339,19 +339,21 @@ def _update_command_executor(self, keep_alive: bool) -> None:
 
         logger.debug('Updated request endpoint to %s', executor)
 
-        # Override command executor.
+        # Keep the existing executor, including custom constructor arguments and state.
         # The client configuration given by a user, e.g. the read timeout, the proxy or
         # the authentication credentials, must be kept as-is. Only the endpoint the client
         # talks to changes, thus a copy of the current configuration is reused instead of
         # building a brand-new one out of the endpoint URL.
         client_config = copy.copy(self.command_executor.client_config)
         client_config.remote_server_addr = executor
         client_config.keep_alive = keep_alive
-        if isinstance(self.command_executor, AppiumConnection):
-            self.command_executor = type(self.command_executor)(client_config=client_config)
-        else:
-            self.command_executor = RemoteConnection(client_config=client_config)
-        self._add_commands()
+        if hasattr(self.command_executor, '_conn'):
+            self.command_executor._conn.clear()
+        self.command_executor._client_config = client_config
+        # Proxy selection (including no_proxy) can change with the endpoint.
+        self.command_executor._proxy_url = client_config.get_proxy_url()
+        if keep_alive:
+            self.command_executor._conn = self.command_executor._get_connection_manager()
 
     # https://github.com/SeleniumHQ/selenium/blob/06fdf2966df6bca47c0ae45e8201cd30db9b9a49/py/selenium/webdriver/remote/webdriver.py#L277
     # noinspection PyAttributeOutsideInit

test/unit/webdriver/webdriver_test.py

@@ -349,8 +349,71 @@ def test_create_session_register_uridirect_ignores_empty_values(self):
 
         assert driver.command_executor._client_config.remote_server_addr == SERVER_URL_BASE
 
+    @pytest.mark.parametrize('prefix', ['', 'appium:'])
+    @pytest.mark.parametrize('path', ['', '/', None])
     @httpretty.activate
-    def test_create_session_register_uridirect_with_custom_connection(self):
+    def test_direct_connect_root_path(self, prefix, path):
+        caps = {
+            f'{prefix}directConnectProtocol': 'http',
+            f'{prefix}directConnectHost': 'localhost2',
+            f'{prefix}directConnectPort': 4800,
+            f'{prefix}directConnectPath': path,
+        }
+        # An explicit empty prefixed path must win over a legacy path.
+        if prefix:
+            caps['directConnectPath'] = '/legacy' if path is not None else None
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body=json.dumps({'sessionId': 'session-id', 'capabilities': caps}),
+        )
+        endpoint = SERVER_URL_BASE if path is None else f'http://localhost2:4800{path}'
+        httpretty.register_uri(
+            httpretty.GET,
+            f'{endpoint}/session/session-id/contexts',
+            body=json.dumps({'value': ['NATIVE_APP']}),
+        )
+        driver = webdriver.Remote(
+            options=UiAutomator2Options(),
+            client_config=AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True),
+        )
+        assert driver.command_executor.client_config.remote_server_addr == endpoint
+        assert driver.contexts == ['NATIVE_APP']
+
+    @pytest.mark.parametrize('keep_alive', [True, False])
+    def test_direct_connect_rebuilds_proxy_pool(self, monkeypatch, keep_alive):
+        monkeypatch.setenv('http_proxy', 'http://proxy:8080')
+        monkeypatch.setenv('no_proxy', 'grid')
+        config = AppiumClientConfig(remote_server_addr='http://grid:4444', keep_alive=keep_alive)
+        connection = AppiumConnection(client_config=config)
+        assert connection._proxy_url is None
+        driver = WebDriver.__new__(WebDriver)
+        driver.command_executor = connection
+        driver.caps = {
+            'directConnectProtocol': 'http',
+            'directConnectHost': 'node',
+            'directConnectPort': 4723,
+            'directConnectPath': '',
+        }
+        with patch.object(connection, '_get_connection_manager', wraps=connection._get_connection_manager) as manager:
+            if keep_alive:
+                old_pool = connection._conn
+                with patch.object(old_pool, 'clear', wraps=old_pool.clear) as clear:
+                    driver._update_command_executor(keep_alive)
+                clear.assert_called_once()
+                assert connection._conn is not old_pool
+                assert isinstance(connection._conn, urllib3.ProxyManager)
+                manager.assert_called_once()
+            else:
+                driver._update_command_executor(keep_alive)
+                manager.assert_not_called()
+        assert connection._proxy_url == 'http://proxy:8080'
+        assert config.remote_server_addr == 'http://grid:4444'
+        connection.close()
+
+    @pytest.mark.parametrize('keep_alive', [True, False])
+    @httpretty.activate
+    def test_create_session_register_uridirect_with_custom_connection(self, keep_alive):
         httpretty.register_uri(
             httpretty.POST,
             f'{SERVER_URL_BASE}/session',
@@ -369,22 +432,38 @@ def test_create_session_register_uridirect_with_custom_connection(self):
         )
 
         class CustomAppiumConnection(AppiumConnection):
-            pass
+            def __init__(self, token, client_config):
+                super().__init__(client_config=client_config)
+                self.token = token
 
         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)
+        client_config = AppiumClientConfig(remote_server_addr=SERVER_URL_BASE, direct_connection=True, keep_alive=keep_alive)
+        custom_connection = CustomAppiumConnection('custom-token', client_config=client_config)
+        state = object()
+        custom_connection.state = state
+        custom_connection._commands = custom_connection._commands.copy()
+        custom_connection.add_command('custom_command', 'GET', '/session/$sessionId/custom')
+        httpretty.register_uri(
+            httpretty.GET,
+            'http://localhost2:4800/special/path/wd/hub/session/session-id/custom',
+            body=json.dumps({'value': 'custom-result'}),
+        )
         driver = webdriver.Remote(
             custom_connection,
             options=UiAutomator2Options().load_capabilities(desired_caps),
         )
 
-        assert isinstance(driver.command_executor, CustomAppiumConnection)
+        assert driver.command_executor is custom_connection
+        assert custom_connection.token == 'custom-token'
+        assert custom_connection.state is state
+        assert client_config.remote_server_addr == SERVER_URL_BASE
+        assert driver.command_executor.client_config.keep_alive is keep_alive
+        assert driver.execute('custom_command')['value'] == 'custom-result'
         assert driver.command_executor._client_config.remote_server_addr == 'http://localhost2:4800/special/path/wd/hub'
 
     @httpretty.activate