chore: better engine cleanup

googleapis · Jan 13, 2025 · 76795b9 · 76795b9
1 parent 8f47a73
commit 76795b9
Show file tree

Hide file tree

Showing 11 changed files with 21 additions and 0 deletions.
diff --git a/tests/test_async_chat_store.py b/tests/test_async_chat_store.py
@@ -98,6 +98,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def chat_store(self, async_engine):

diff --git a/tests/test_async_document_store.py b/tests/test_async_document_store.py
@@ -97,6 +97,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def doc_store(self, async_engine):

diff --git a/tests/test_async_index_store.py b/tests/test_async_index_store.py
@@ -100,6 +100,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def index_store(self, async_engine):

diff --git a/tests/test_async_vector_store.py b/tests/test_async_vector_store.py
@@ -118,6 +118,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE}"')
         await aexecute(engine, f'DROP TABLE IF EXISTS "{DEFAULT_TABLE_CUSTOM_VS}"')
         await engine.close()
+        await engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):

diff --git a/tests/test_async_vector_store_index.py b/tests/test_async_vector_store_index.py
@@ -104,6 +104,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         yield engine
         await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE}")
         await engine.close()
+        await engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):

diff --git a/tests/test_chat_store.py b/tests/test_chat_store.py
@@ -103,6 +103,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def async_chat_store(self, async_engine):
@@ -272,6 +273,7 @@ async def sync_engine(
         yield sync_engine
 
         await sync_engine.close()
+        await sync_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def sync_chat_store(self, sync_engine):

diff --git a/tests/test_document_store.py b/tests/test_document_store.py
@@ -108,6 +108,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def doc_store(self, async_engine):
@@ -400,6 +401,7 @@ async def sync_engine(
         yield sync_engine
 
         await sync_engine.close()
+        await sync_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def sync_doc_store(self, sync_engine):

diff --git a/tests/test_engine.py b/tests/test_engine.py
@@ -122,6 +122,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         await aexecute(engine, f'DROP TABLE "{DEFAULT_IS_TABLE}"')
         await aexecute(engine, f'DROP TABLE "{DEFAULT_CS_TABLE}"')
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_init_with_constructor(
         self,
@@ -178,6 +179,7 @@ async def test_password(
         await aexecute(engine, "SELECT 1")
         AlloyDBEngine._connector = None
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_missing_user_or_password(
         self,
@@ -304,6 +306,7 @@ async def test_iam_account_override(
         assert engine
         await aexecute(engine, "SELECT 1")
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_init_document_store(self, engine):
         await engine.ainit_doc_store_table(
@@ -440,6 +443,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         await aexecute(engine, f'DROP TABLE "{DEFAULT_VS_TABLE_SYNC}"')
         await aexecute(engine, f'DROP TABLE "{DEFAULT_CS_TABLE_SYNC}"')
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_password(
         self,
@@ -465,6 +469,7 @@ async def test_password(
         await aexecute(engine, "SELECT 1")
         AlloyDBEngine._connector = None
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_engine_constructor_key(
         self,
@@ -494,6 +499,7 @@ async def test_iam_account_override(
         assert engine
         await aexecute(engine, "SELECT 1")
         await engine.close()
+        await engine._connector.close_async()
 
     async def test_init_document_store(self, engine):
         engine.init_doc_store_table(

diff --git a/tests/test_index_store.py b/tests/test_index_store.py
@@ -108,6 +108,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def index_store(self, async_engine):
@@ -236,6 +237,7 @@ async def async_engine(
         yield async_engine
 
         await async_engine.close()
+        await async_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def index_store(self, async_engine):

diff --git a/tests/test_vector_store.py b/tests/test_vector_store.py
@@ -124,6 +124,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
 
         await aexecute(sync_engine, f'DROP TABLE "{DEFAULT_TABLE}"')
         await sync_engine.close()
+        await sync_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):
@@ -504,6 +505,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
 
         await aexecute(sync_engine, f'DROP TABLE "{DEFAULT_TABLE}"')
         await sync_engine.close()
+        await sync_engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):

diff --git a/tests/test_vector_store_index.py b/tests/test_vector_store_index.py
@@ -116,6 +116,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         yield engine
         await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE}")
         await engine.close()
+        await engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):
@@ -292,6 +293,7 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
         yield engine
         await aexecute(engine, f"DROP TABLE IF EXISTS {DEFAULT_TABLE_ASYNC}")
         await engine.close()
+        await engine._connector.close_async()
 
     @pytest_asyncio.fixture(scope="class")
     async def vs(self, engine):