chore: improve task and test cleanup (#277)

Author: jackwotherspoon

google/cloud/sql/connector/connector.py

@@ -150,6 +150,23 @@ def connect(
             icm.force_refresh()
             raise (e)
 
+    async def _close(self) -> None:
+        """Helper function to close InstanceConnectionManagers' tasks."""
+        await asyncio.gather(*[icm.close() for icm in self._instances.values()])
+
+    def __del__(self) -> None:
+        """Deconstructor to make sure InstanceConnectionManagers are closed
+        and tasks have finished to have a graceful exit.
+        """
+        logger.debug("Entering deconstructor")
+
+        deconstruct_future = asyncio.run_coroutine_threadsafe(
+            self._close(), loop=self._loop
+        )
+        # Will attempt to safely shut down tasks for 5s
+        deconstruct_future.result(timeout=5)
+        logger.debug("Finished deconstructing")
+
 
 def connect(instance_connection_string: str, driver: str, **kwargs: Any) -> Any:
     """Uses a Connector object with default settings and returns a database

google/cloud/sql/connector/instance_connection_manager.py

@@ -280,30 +280,6 @@ async def _async_init() -> None:
         init_future = asyncio.run_coroutine_threadsafe(_async_init(), self._loop)
         init_future.result()
 
-    def __del__(self) -> None:
-        """Deconstructor to make sure ClientSession is closed and tasks have
-        finished to have a graceful exit.
-        """
-        logger.debug("Entering deconstructor")
-
-        async def _deconstruct() -> None:
-            if isinstance(self._current, asyncio.Task):
-                logger.debug("Waiting for _current to be cancelled")
-                self._current.cancel()
-            if isinstance(self._next, asyncio.Task):
-                logger.debug("Waiting for _next to be cancelled")
-                self._next.cancel()
-            if not self._client_session.closed:
-                logger.debug("Waiting for _client_session to close")
-                await self._client_session.close()
-
-        deconstruct_future = asyncio.run_coroutine_threadsafe(
-            _deconstruct(), loop=self._loop
-        )
-        # Will attempt to safely shut down tasks for 5s
-        deconstruct_future.result(timeout=5)
-        logger.debug("Finished deconstructing")
-
     async def _get_instance_data(self) -> InstanceMetadata:
         """Asynchronous function that takes in the futures for the ephemeral certificate
         and the instance metadata and generates an OpenSSL context object.
@@ -687,3 +663,14 @@ def _connect_with_pytds(
         return pytds.connect(
             ip_address, database=db, user=user, password=passwd, sock=sock, **kwargs
         )
+
+    async def close(self) -> None:
+        """Cleanup function to make sure ClientSession is closed and tasks have
+        finished to have a graceful exit.
+        """
+        logger.debug("Waiting for _current to be cancelled")
+        self._current.cancel()
+        logger.debug("Waiting for _next to be cancelled")
+        self._next.cancel()
+        logger.debug("Waiting for _client_session to close")
+        await self._client_session.close()

tests/unit/test_connector.py

@@ -101,6 +101,8 @@ def test_connect_enable_iam_auth_error() -> None:
             "pg8000",
             enable_iam_auth=True,
         )
+        # remove mock_icm to avoid destructor warnings
+        default_connector._instances = {}
 
 
 def test_default_Connector_Init() -> None:

tests/unit/test_rate_limiter.py

@@ -60,6 +60,10 @@ async def increment() -> None:
     assert len(done) == 4
     assert len(pending) == 6
 
+    # cleanup pending tasks
+    for task in pending:
+        task.cancel()
+
 
 @pytest.mark.asyncio
 async def test_rate_limiter_completes_all_tasks(