|
13 | 13 | # limitations under the License. |
14 | 14 |
|
15 | 15 | import uuid |
16 | | -from typing import TYPE_CHECKING, Any, Dict, Optional, Union |
| 16 | +from typing import TYPE_CHECKING, Any, Dict |
17 | 17 |
|
18 | | -import urllib3 |
19 | 18 | from selenium.webdriver.remote.remote_connection import RemoteConnection |
20 | 19 |
|
21 | 20 | from appium.common.helper import library_version |
|
26 | 25 |
|
27 | 26 | PREFIX_HEADER = 'appium/' |
28 | 27 |
|
| 28 | +HEADER_IDEMOTENCY_KEY = 'X-Idempotency-Key' |
29 | 29 |
|
30 | | -class AppiumConnection(RemoteConnection): |
31 | | - _proxy_url: Optional[str] |
32 | | - |
33 | | - def __init__( |
34 | | - self, |
35 | | - remote_server_addr: str, |
36 | | - keep_alive: bool = False, |
37 | | - ignore_proxy: Optional[bool] = False, |
38 | | - init_args_for_pool_manager: Union[Dict[str, Any], None] = None, |
39 | | - ): |
40 | | - # Need to call before super().__init__ in order to pass arguments for the pool manager in the super. |
41 | | - self._init_args_for_pool_manager = init_args_for_pool_manager or {} |
42 | | - |
43 | | - super().__init__(remote_server_addr, keep_alive=keep_alive, ignore_proxy=ignore_proxy) |
44 | | - |
45 | | - def _get_connection_manager(self) -> Union[urllib3.PoolManager, urllib3.ProxyManager]: |
46 | | - # https://github.com/SeleniumHQ/selenium/blob/0e0194b0e52a34e7df4b841f1ed74506beea5c3e/py/selenium/webdriver/remote/remote_connection.py#L134 |
47 | | - pool_manager_init_args = {'timeout': self.get_timeout()} |
48 | | - |
49 | | - if self._ca_certs: |
50 | | - pool_manager_init_args['cert_reqs'] = 'CERT_REQUIRED' |
51 | | - pool_manager_init_args['ca_certs'] = self._ca_certs |
52 | | - else: |
53 | | - # This line is necessary to disable certificate verification |
54 | | - pool_manager_init_args['cert_reqs'] = 'CERT_NONE' |
55 | 30 |
|
56 | | - pool_manager_init_args.update(self._init_args_for_pool_manager) |
| 31 | +def _get_new_headers(key: str, headers: Dict[str, str]) -> Dict[str, str]: |
| 32 | + """Return a new dictionary of heafers without the given key. |
| 33 | + The key match is case-insensitive.""" |
| 34 | + lower_key = key.lower() |
| 35 | + return {k: v for k, v in headers.items() if k.lower() != lower_key} |
| 36 | + |
57 | 37 |
|
58 | | - if self._proxy_url: |
59 | | - if self._proxy_url.lower().startswith('sock'): |
60 | | - from urllib3.contrib.socks import SOCKSProxyManager |
| 38 | +class AppiumConnection(RemoteConnection): |
| 39 | + """ |
| 40 | + A subclass of selenium.webdriver.remote.remote_connection.Remoteconnection. |
61 | 41 |
|
62 | | - return SOCKSProxyManager(self._proxy_url, **pool_manager_init_args) |
63 | | - if self._identify_http_proxy_auth(): |
64 | | - self._proxy_url, self._basic_proxy_auth = self._separate_http_proxy_auth() |
65 | | - pool_manager_init_args['proxy_headers'] = urllib3.make_headers(proxy_basic_auth=self._basic_proxy_auth) |
66 | | - return urllib3.ProxyManager(self._proxy_url, **pool_manager_init_args) |
| 42 | + The changes are: |
| 43 | + - The default user agent |
| 44 | + - Adds 'X-Idempotency-Key' header in a new session request to avoid proceeding |
| 45 | + the same request multiple times in the Appium server side. |
| 46 | + - https://github.com/appium/appium-base-driver/pull/400 |
| 47 | + """ |
67 | 48 |
|
68 | | - return urllib3.PoolManager(**pool_manager_init_args) |
| 49 | + user_agent = f'{PREFIX_HEADER}{library_version()} ({RemoteConnection.user_agent})' |
| 50 | + extra_headers = {} |
69 | 51 |
|
70 | 52 | @classmethod |
71 | 53 | def get_remote_connection_headers(cls, parsed_url: 'ParseResult', keep_alive: bool = True) -> Dict[str, Any]: |
72 | | - """Override get_remote_connection_headers in RemoteConnection""" |
73 | | - headers = RemoteConnection.get_remote_connection_headers(parsed_url, keep_alive=keep_alive) |
74 | | - # e.g. appium/0.49 (selenium/3.141.0 (python linux)) |
75 | | - headers['User-Agent'] = f'{PREFIX_HEADER}{library_version()} ({headers["User-Agent"]})' |
| 54 | + """Override get_remote_connection_headers in RemoteConnection to control the extra headers. |
| 55 | + This method will be used in sending a request method in this class. |
| 56 | + """ |
| 57 | + |
76 | 58 | if parsed_url.path.endswith('/session'): |
77 | 59 | # https://github.com/appium/appium-base-driver/pull/400 |
78 | | - headers['X-Idempotency-Key'] = str(uuid.uuid4()) |
| 60 | + cls.extra_headers[HEADER_IDEMOTENCY_KEY] = str(uuid.uuid4()) |
| 61 | + else: |
| 62 | + cls.extra_headers = _get_new_headers(HEADER_IDEMOTENCY_KEY, cls.extra_headers) |
79 | 63 |
|
80 | | - return headers |
| 64 | + return {**super().get_remote_connection_headers(parsed_url, keep_alive=keep_alive), **cls.extra_headers} |
0 commit comments