From 76795b910948ae6be72977ae60b757dee8bcec5c Mon Sep 17 00:00:00 2001 From: Vishwaraj Anand Date: Mon, 13 Jan 2025 18:30:03 +0000 Subject: [PATCH] chore: better engine cleanup --- tests/test_async_chat_store.py | 1 + tests/test_async_document_store.py | 1 + tests/test_async_index_store.py | 1 + tests/test_async_vector_store.py | 1 + tests/test_async_vector_store_index.py | 1 + tests/test_chat_store.py | 2 ++ tests/test_document_store.py | 2 ++ tests/test_engine.py | 6 ++++++ tests/test_index_store.py | 2 ++ tests/test_vector_store.py | 2 ++ tests/test_vector_store_index.py | 2 ++ 11 files changed, 21 insertions(+) diff --git a/tests/test_async_chat_store.py b/tests/test_async_chat_store.py index 70397e4..3ee43ad 100644 --- 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 index c978f94..d6f483f 100644 --- 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 index 09f532d..6accfa8 100644 --- 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 index e99ac51..20af41c 100644 --- 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 index f7373c7..2440866 100644 --- 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 index 087c880..17f5cc3 100644 --- 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 index 44533e4..840d2eb 100644 --- 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 index fd168e1..9a08f62 100644 --- 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 index a5bc09e..5b905dd 100644 --- 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 index 96a7bca..50fb5f9 100644 --- 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 index 3920929..c0a5f8a 100644 --- 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):