Skip to content

Commit

Permalink
Fix failing tests (#392)
Browse files Browse the repository at this point in the history
Signed-off-by: Levko Kravets <levko.ne@gmail.com>
  • Loading branch information
kravets-levko authored May 28, 2024
1 parent d63b71b commit a6e9b11
Show file tree
Hide file tree
Showing 2 changed files with 6 additions and 5 deletions.
6 changes: 3 additions & 3 deletions src/databricks/sql/auth/thrift_http_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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()}
Expand All @@ -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" % (
Expand Down
5 changes: 3 additions & 2 deletions tests/unit/test_thrift_backend.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down

0 comments on commit a6e9b11

Please sign in to comment.