Skip to content

Commit

Permalink
fix: clean up HTTP session and pool during tear down phase
Browse files Browse the repository at this point in the history
* Add unit tests for the change
* Fix the unittest to test on the correct class
* Make linter happy
  • Loading branch information
lidizheng committed Apr 6, 2022
1 parent dc040f5 commit 76bf214
Show file tree
Hide file tree
Showing 4 changed files with 26 additions and 0 deletions.
4 changes: 4 additions & 0 deletions google/auth/transport/requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -149,6 +149,10 @@ def __init__(self, session=None):

self.session = session

def __del__(self):
if hasattr(self, "session") and self.session is not None:
self.session.close()

def __call__(
self,
url,
Expand Down
4 changes: 4 additions & 0 deletions google/auth/transport/urllib3.py
Original file line number Diff line number Diff line change
Expand Up @@ -427,6 +427,10 @@ def __exit__(self, exc_type, exc_val, exc_tb):
"""Proxy to ``self.http``."""
return self.http.__exit__(exc_type, exc_val, exc_tb)

def __del__(self):
if hasattr(self, "http") and self.http is not None:
self.http.clear()

@property
def headers(self):
"""Proxy to ``self.http``."""
Expand Down
6 changes: 6 additions & 0 deletions tests/transport/test_requests.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,12 @@ def test_timeout(self):

assert http.request.call_args[1]["timeout"] == 5

def test_session_closed_on_del(self):
http = mock.create_autospec(requests.Session, instance=True)
request = google.auth.transport.requests.Request(http)
request.__del__()
http.close.assert_called_with()


class TestTimeoutGuard(object):
def make_guard(self, *args, **kwargs):
Expand Down
12 changes: 12 additions & 0 deletions tests/transport/test_urllib3.py
Original file line number Diff line number Diff line change
Expand Up @@ -305,3 +305,15 @@ def test_configure_mtls_channel_without_client_cert_env(
is_mtls = authed_http.configure_mtls_channel(callback)
assert not is_mtls
get_client_cert_and_key.assert_not_called()

def test_clear_pool_on_del(self):
http = mock.create_autospec(urllib3.PoolManager)
authed_http = google.auth.transport.urllib3.AuthorizedHttp(
mock.sentinel.credentials, http=http
)
authed_http.__del__()
http.clear.assert_called_with()

authed_http.http = None
authed_http.__del__()
# Expect it to not crash

0 comments on commit 76bf214

Please sign in to comment.