feat: respect the given executor (#844)

KazuCocoa · web-flow · commit 2c92c04b3d47 · 2023-04-01T13:27:12.000-07:00
* feat: can provide a custom connection

* add tests

* tweak tests

* lint

* tweak

* add comment

* fix lint

* tweak

* add test

* tweak review
diff --git a/README.md b/README.md
@@ -115,7 +115,7 @@ options = UiAutomator2Options()
 options.platformVersion = '10'
 options.udid = '123456789ABC'
 options.app = PATH('../../../apps/test-app.apk')
-# Appium1 points to http://127.0.0.1:4723/wd/hub by default 
+# Appium1 points to http://127.0.0.1:4723/wd/hub by default
 self.driver = webdriver.Remote('http://127.0.0.1:4723', options=options)
 el = self.driver.find_element(by=AppiumBy.ACCESSIBILITY_ID, value='item')
 el.click()
@@ -134,7 +134,7 @@ options = XCUITestOptions()
 options.platformVersion = '13.4'
 options.udid = '123456789ABC'
 options.app = PATH('../../apps/UICatalog.app.zip')
-# Appium1 points to http://127.0.0.1:4723/wd/hub by default 
+# Appium1 points to http://127.0.0.1:4723/wd/hub by default
 self.driver = webdriver.Remote('http://127.0.0.1:4723', options=options)
 el = self.driver.find_element(by=AppiumBy.ACCESSIBILITY_ID, value='item')
 el.click()
@@ -158,7 +158,7 @@ from appium import webdriver
 # instead: https://github.com/appium/python-client/pull/720
 from appium.options.ios import XCUITestOptions
 
-# load_capabilities API could be used to 
+# load_capabilities API could be used to
 # load options mapping stored in a dictionary
 options = XCUITestOptions().load_capabilities({
     'platformVersion': '13.4',
@@ -167,9 +167,9 @@ options = XCUITestOptions().load_capabilities({
 })
 
 self.driver = webdriver.Remote(
-    # Appium1 points to http://127.0.0.1:4723/wd/hub by default 
-    'http://127.0.0.1:4723', 
-    options=options, 
+    # Appium1 points to http://127.0.0.1:4723/wd/hub by default
+    'http://127.0.0.1:4723',
+    options=options,
     direct_connection=True
 )
 ```
@@ -192,10 +192,62 @@ options.automation_name = 'safari'
 # calls to it could be chained
 options.set_capability('browser_name', 'safari')
 
-# Appium1 points to http://127.0.0.1:4723/wd/hub by default 
+# Appium1 points to http://127.0.0.1:4723/wd/hub by default
 self.driver = webdriver.Remote('http://127.0.0.1:4723', options=options, strict_ssl=False)
 ```
 
+## Set custom `AppiumConnection`
+
+The first argument of `webdriver.Remote` can set an arbitrary command executor for you.
+
+1. Set init arguments for the pool manager Appium Python client uses to manage http requests.
+
+```python
+from appium import webdriver
+from appium.options.ios import XCUITestOptions
+
+import urllib3
+from appium.webdriver.appium_connection import AppiumConnection
+
+# Retry connection error up to 3 times.
+init_args_for_pool_manage = {
+    'retries': urllib3.util.retry.Retry(total=3, connect=3, read=False)
+}
+appium_executor = AppiumConnection(remote_server_addr='http://127.0.0.1:4723',
+    init_args_for_pool_manage=init_args_for_pool_manage)
+
+options = XCUITestOptions()
+options.platformVersion = '13.4'
+options.udid = '123456789ABC'
+options.app = PATH('../../apps/UICatalog.app.zip')
+driver = webdriver.Remote(appium_executor, options=options)
+```
+
+
+2. Define a subclass of `AppiumConnection`
+
+```python
+from appium import webdriver
+from appium.options.ios import XCUITestOptions
+
+from appium.webdriver.appium_connection import AppiumConnection
+
+class CustomAppiumConnection(AppiumConnection):
+    # Can add your own methods for the custom class
+    pass
+
+custom_executor = CustomAppiumConnection(remote_server_addr='http://127.0.0.1:4723')
+
+options = XCUITestOptions().load_capabilities({
+    'platformVersion': '13.4',
+    'deviceName': 'iPhone Simulator',
+    'app': PATH('../../apps/UICatalog.app.zip'),
+})
+driver = webdriver.Remote(custom_executor, options=options)
+
+```
+
+
 ## Documentation
 
 - https://appium.github.io/python-client-sphinx/ is detailed documentation
diff --git a/appium/webdriver/appium_connection.py b/appium/webdriver/appium_connection.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 
 import uuid
-from typing import TYPE_CHECKING, Any, Dict, Union
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 
 import urllib3
 from selenium.webdriver.remote.remote_connection import RemoteConnection
@@ -25,17 +25,32 @@
 
 
 class AppiumConnection(RemoteConnection):
+    def __init__(
+        self,
+        remote_server_addr: str,
+        keep_alive: bool = False,
+        ignore_proxy: Optional[bool] = False,
+        init_args_for_pool_manager: Union[Dict[str, Any], None] = None,
+    ):
+
+        # Need to call before super().__init__ in order to pass arguments for the pool manager in the super.
+        self._init_args_for_pool_manager = init_args_for_pool_manager or {}
+
+        super().__init__(remote_server_addr, keep_alive=keep_alive, ignore_proxy=ignore_proxy)
+
     def _get_connection_manager(self) -> Union[urllib3.PoolManager, urllib3.ProxyManager]:
         # https://github.com/SeleniumHQ/selenium/blob/0e0194b0e52a34e7df4b841f1ed74506beea5c3e/py/selenium/webdriver/remote/remote_connection.py#L134
         pool_manager_init_args = {'timeout': self._timeout}
-        # pylint: disable=E1101
+
         if self._ca_certs:
             pool_manager_init_args['cert_reqs'] = 'CERT_REQUIRED'
             pool_manager_init_args['ca_certs'] = self._ca_certs
         else:
             # This line is necessary to disable certificate verification
             pool_manager_init_args['cert_reqs'] = 'CERT_NONE'
 
+        pool_manager_init_args.update(self._init_args_for_pool_manager)
+
         return (
             urllib3.PoolManager(**pool_manager_init_args)
             if self._proxy_url is None
diff --git a/appium/webdriver/webdriver.py b/appium/webdriver/webdriver.py
@@ -204,7 +204,7 @@ class WebDriver(
 ):
     def __init__(
         self,
-        command_executor: str = 'http://127.0.0.1:4444/wd/hub',
+        command_executor: Union[str, AppiumConnection] = 'http://127.0.0.1:4444/wd/hub',
         desired_capabilities: Optional[Dict] = None,
         browser_profile: Union[str, None] = None,
         proxy: Union[str, None] = None,
@@ -228,8 +228,11 @@ def __init__(
             AppiumConnection.set_certificate_bundle_path(None)
             urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
+        if isinstance(command_executor, str):
+            command_executor = AppiumConnection(command_executor, keep_alive=keep_alive)
+
         super().__init__(
-            command_executor=AppiumConnection(command_executor, keep_alive=keep_alive),
+            command_executor=command_executor,
             desired_capabilities=desired_capabilities,
             browser_profile=browser_profile,
             proxy=proxy,
diff --git a/test/unit/webdriver/webdriver_test.py b/test/unit/webdriver/webdriver_test.py
@@ -15,11 +15,13 @@
 import json
 
 import httpretty
+import urllib3
 from mock import patch
 
 from appium import version as appium_version
 from appium import webdriver
 from appium.options.android import UiAutomator2Options
+from appium.webdriver.appium_connection import AppiumConnection
 from appium.webdriver.webdriver import ExtensionBase, WebDriver
 from test.helpers.constants import SERVER_URL_BASE
 from test.unit.helper.test_helper import (
@@ -298,6 +300,95 @@ def add_command(self):
         assert result == {}
         driver.delete_extensions()
 
+    @httpretty.activate
+    def test_create_session_with_custom_connection(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body='{ "value": {"sessionId": "session-id", "capabilities": {"deviceName": "Android Emulator"}} }',
+        )
+
+        desired_caps = {
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+        }
+
+        class CustomAppiumConnection(AppiumConnection):
+            # To explicitly check if the given executor is used
+            pass
+
+        init_args_for_pool_manager = {'retries': urllib3.util.retry.Retry(total=3, connect=3, read=False)}
+        custom_appium_connection = CustomAppiumConnection(
+            remote_server_addr=SERVER_URL_BASE, init_args_for_pool_manager=init_args_for_pool_manager
+        )
+
+        driver = webdriver.Remote(
+            custom_appium_connection, options=UiAutomator2Options().load_capabilities(desired_caps)
+        )
+
+        request = httpretty.HTTPretty.latest_requests[0]
+        assert request.headers['content-type'] == 'application/json;charset=UTF-8'
+        assert 'appium/python {} (selenium'.format(appium_version.version) in request.headers['user-agent']
+
+        request_json = json.loads(httpretty.HTTPretty.latest_requests[0].body.decode('utf-8'))
+        assert request_json.get('capabilities') is not None
+        assert request_json['capabilities']['alwaysMatch'] == {
+            'platformName': 'Android',
+            'appium:deviceName': 'Android Emulator',
+            'appium:app': 'path/to/app',
+            'appium:automationName': 'UIAutomator2',
+        }
+        assert request_json.get('desiredCapabilities') is None
+        assert driver.session_id == 'session-id'
+
+        assert isinstance(driver.command_executor, CustomAppiumConnection)
+
+    @httpretty.activate
+    def test_create_session_with_custom_connection_with_keepalive(self):
+        httpretty.register_uri(
+            httpretty.POST,
+            f'{SERVER_URL_BASE}/session',
+            body='{ "value": {"sessionId": "session-id", "capabilities": {"deviceName": "Android Emulator"}} }',
+        )
+
+        desired_caps = {
+            'deviceName': 'Android Emulator',
+            'app': 'path/to/app',
+        }
+
+        class CustomAppiumConnection(AppiumConnection):
+            # To explicitly check if the given executor is used
+            pass
+
+        init_args_for_pool_manager = {'retries': urllib3.util.retry.Retry(total=3, connect=3, read=False)}
+        custom_appium_connection = CustomAppiumConnection(
+            # keep alive has different route to set init args for the pool manager
+            keep_alive=True,
+            remote_server_addr=SERVER_URL_BASE,
+            init_args_for_pool_manager=init_args_for_pool_manager,
+        )
+
+        driver = webdriver.Remote(
+            custom_appium_connection, options=UiAutomator2Options().load_capabilities(desired_caps)
+        )
+
+        request = httpretty.HTTPretty.latest_requests[0]
+        assert request.headers['content-type'] == 'application/json;charset=UTF-8'
+        assert 'appium/python {} (selenium'.format(appium_version.version) in request.headers['user-agent']
+
+        request_json = json.loads(httpretty.HTTPretty.latest_requests[0].body.decode('utf-8'))
+        assert request_json.get('capabilities') is not None
+        assert request_json['capabilities']['alwaysMatch'] == {
+            'platformName': 'Android',
+            'appium:deviceName': 'Android Emulator',
+            'appium:app': 'path/to/app',
+            'appium:automationName': 'UIAutomator2',
+        }
+        assert request_json.get('desiredCapabilities') is None
+        assert driver.session_id == 'session-id'
+
+        assert isinstance(driver.command_executor, CustomAppiumConnection)
+
 
 class SubWebDriver(WebDriver):
     def __init__(self, command_executor, desired_capabilities=None, direct_connection=False, options=None):
