fix: clean up HTTP session and pool during tear down phase (googleapis#1007)

lidizheng · arithmetic1728 · web-flow · commit d05737624528 · 2022-04-05T22:17:54.000-07:00
* Add unit tests for the change
* Fix the unittest to test on the correct class
* Make linter happy

Co-authored-by: arithmetic1728 &lt;58957152+arithmetic1728@users.noreply.github.com&gt;
diff --git a/google/auth/transport/requests.py b/google/auth/transport/requests.py
@@ -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,
diff --git a/google/auth/transport/urllib3.py b/google/auth/transport/urllib3.py
@@ -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``."""
diff --git a/tests/transport/test_requests.py b/tests/transport/test_requests.py
@@ -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):
diff --git a/tests/transport/test_urllib3.py b/tests/transport/test_urllib3.py
@@ -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
