From 7c302be174d5c36444ee97a8712663afc8504d69 Mon Sep 17 00:00:00 2001
From: Shubham Singh <41840111+singh811@users.noreply.github.com>
Date: Thu, 20 Jun 2019 06:56:21 +0530
Subject: [PATCH] [py] Added socksVersion field in proxy class (#6983)

---
 py/selenium/webdriver/common/proxy.py         | 24 +++++++++++++++++++
 .../selenium/webdriver/common/proxy_tests.py  |  4 ++++
 2 files changed, 28 insertions(+)

diff --git a/py/selenium/webdriver/common/proxy.py b/py/selenium/webdriver/common/proxy.py
index ccd362af62733..2dd11a7b162ad 100644
--- a/py/selenium/webdriver/common/proxy.py
+++ b/py/selenium/webdriver/common/proxy.py
@@ -77,6 +77,7 @@ class Proxy(object):
     socksProxy = ''
     socksUsername = ''
     socksPassword = ''
+    socksVersion = None
 
     def __init__(self, raw=None):
         """
@@ -106,6 +107,8 @@ def __init__(self, raw=None):
                 self.socks_username = raw['socksUsername']
             if 'socksPassword' in raw and raw['socksPassword'] is not None:
                 self.socks_password = raw['socksPassword']
+            if 'socksVersion' in raw and raw['socksVersion'] is not None:
+                self.socks_version = raw['socksVersion']
 
     @property
     def proxy_type(self):
@@ -300,6 +303,25 @@ def socks_password(self, value):
         self.proxyType = ProxyType.MANUAL
         self.socksPassword = value
 
+    @property
+    def socks_version(self):
+        """
+        Returns socks proxy version setting.
+        """
+        return self.socksVersion
+
+    @socks_version.setter
+    def socks_version(self, value):
+        """
+        Sets socks proxy version setting.
+
+        :Args:
+         - value: The socks proxy version value.
+        """
+        self._verify_proxy_type_compatibility(ProxyType.MANUAL)
+        self.proxyType = ProxyType.MANUAL
+        self.socksVersion = value
+
     def _verify_proxy_type_compatibility(self, compatibleProxy):
         if self.proxyType != ProxyType.UNSPECIFIED and self.proxyType != compatibleProxy:
             raise Exception(" Specified proxy type (%s) not compatible with current setting (%s)" % (compatibleProxy, self.proxyType))
@@ -331,4 +353,6 @@ def add_to_capabilities(self, capabilities):
             proxy_caps['socksUsername'] = self.socksUsername
         if self.socksPassword:
             proxy_caps['socksPassword'] = self.socksPassword
+        if self.socksVersion:
+            proxy_caps['socksVersion'] = self.socksVersion
         capabilities['proxy'] = proxy_caps
diff --git a/py/test/selenium/webdriver/common/proxy_tests.py b/py/test/selenium/webdriver/common/proxy_tests.py
index 65736c9ec39e4..b351f8336aa40 100755
--- a/py/test/selenium/webdriver/common/proxy_tests.py
+++ b/py/test/selenium/webdriver/common/proxy_tests.py
@@ -28,6 +28,7 @@
     'socksProxy': 'socks.proxy:65555',
     'socksUsername': 'test',
     'socksPassword': 'test',
+    'socksVersion': 5,
 }
 
 PAC_PROXY = {
@@ -48,6 +49,7 @@ def testCanAddManualProxyToDesiredCapabilities():
     proxy.socksProxy = MANUAL_PROXY['socksProxy']
     proxy.socksUsername = MANUAL_PROXY['socksUsername']
     proxy.socksPassword = MANUAL_PROXY['socksPassword']
+    proxy.socksVersion = MANUAL_PROXY['socksVersion']
 
     desired_capabilities = {}
     proxy.add_to_capabilities(desired_capabilities)
@@ -105,6 +107,7 @@ def testCanInitManualProxy():
     assert MANUAL_PROXY['socksProxy'] == proxy.socksProxy
     assert MANUAL_PROXY['socksUsername'] == proxy.socksUsername
     assert MANUAL_PROXY['socksPassword'] == proxy.socksPassword
+    assert MANUAL_PROXY['socksVersion'] == proxy.socksVersion
 
 
 def testCanInitAutodetectProxy():
@@ -131,6 +134,7 @@ def testCanInitEmptyProxy():
     assert '' == proxy.socksPassword
     assert proxy.auto_detect is False
     assert '' == proxy.proxy_autoconfig_url
+    assert proxy.socks_version is None
 
     desired_capabilities = {}
     proxy.add_to_capabilities(desired_capabilities)