From 5879eb34cc784637f82fa040174112d8f6f7f838 Mon Sep 17 00:00:00 2001 From: Levko Kravets Date: Thu, 23 May 2024 16:37:34 +0300 Subject: [PATCH] Fix failing tests Signed-off-by: Levko Kravets --- src/databricks/sql/auth/thrift_http_client.py | 6 +++--- tests/unit/test_thrift_backend.py | 5 +++-- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/databricks/sql/auth/thrift_http_client.py b/src/databricks/sql/auth/thrift_http_client.py index 5db76de97..f7c22a1ed 100644 --- a/src/databricks/sql/auth/thrift_http_client.py +++ b/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" % ( diff --git a/tests/unit/test_thrift_backend.py b/tests/unit/test_thrift_backend.py index ee4fc4b72..4bcf84d2a 100644 --- a/tests/unit/test_thrift_backend.py +++ b/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")