Fix failing tests (#392)

Signed-off-by: Levko Kravets <[email protected]>

Author: kravets-levko

src/databricks/sql/auth/thrift_http_client.py

@@ -78,7 +78,7 @@ def __init__(
             self.proxy_uri: str = proxy
             self.host = parsed.hostname
             self.port = parsed.port
-            self.proxy_auth = self.basic_proxy_auth_header(parsed)
+            self.proxy_auth = self.basic_proxy_auth_headers(parsed)
         else:
             self.realhost = self.realport = self.proxy_auth = None
 
@@ -167,7 +167,7 @@ def flush(self):
         }
 
         if self.using_proxy() and self.scheme == "http" and self.proxy_auth is not None:
-            headers["Proxy-Authorization" : self.proxy_auth]
+            headers.update(self.proxy_auth)
 
         if self.__custom_headers:
             custom_headers = {key: val for key, val in self.__custom_headers.items()}
@@ -190,7 +190,7 @@ def flush(self):
         self.headers = self.__resp.headers
 
     @staticmethod
-    def basic_proxy_auth_header(proxy):
+    def basic_proxy_auth_headers(proxy):
         if proxy is None or not proxy.username:
             return None
         ap = "%s:%s" % (

tests/unit/test_thrift_backend.py

@@ -167,11 +167,12 @@ def test_proxy_headers_are_set(self):
         parsed_proxy = urlparse(fake_proxy_spec)
 
         try:
-            result = THttpClient.basic_proxy_auth_header(parsed_proxy)
+            result = THttpClient.basic_proxy_auth_headers(parsed_proxy)
         except TypeError as e:
             assert False
 
-        assert isinstance(result, type(str()))
+        assert isinstance(result, type(dict()))
+        assert isinstance(result.get('proxy-authorization'), type(str()))
 
     @patch("databricks.sql.auth.thrift_http_client.THttpClient")
     @patch("databricks.sql.thrift_backend.create_default_context")