Improved tests for async SQLAlchemy engines [skip ci]

Author: ankane

tests/test_sqlalchemy.py

@@ -19,11 +19,18 @@
 psycopg2_engine = create_engine('postgresql+psycopg2://localhost/pgvector_python_test')
 pg8000_engine = create_engine(f'postgresql+pg8000://{os.environ["USER"]}@localhost/pgvector_python_test')
 engines = [psycopg2_engine, pg8000_engine]
+async_engines = []
 
 if sqlalchemy_version > 1:
     psycopg_engine = create_engine('postgresql+psycopg://localhost/pgvector_python_test')
     engines.append(psycopg_engine)
 
+    psycopg_async_engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')
+    async_engines.append(psycopg_async_engine)
+
+    asyncpg_engine = create_async_engine('postgresql+asyncpg://localhost/pgvector_python_test')
+    async_engines.append(asyncpg_engine)
+
 setup_engine = engines[0]
 with Session(setup_engine) as session:
     session.execute(text('CREATE EXTENSION IF NOT EXISTS vector'))
@@ -531,14 +538,13 @@ def test_halfvec_array(self, engine):
             assert item.half_embeddings[1].to_list() == [4, 5, 6]
 
 
+@pytest.mark.parametrize('engine', async_engines)
 class TestSqlalchemyAsync:
     def setup_method(self):
         delete_items()
 
     @pytest.mark.asyncio
-    @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
-    async def test_psycopg_async_avg(self):
-        engine = create_async_engine('postgresql+psycopg://localhost/pgvector_python_test')
+    async def test_psycopg_async_avg(self, engine):
         async_session = async_sessionmaker(engine, expire_on_commit=False)
 
         async with async_session() as session:
@@ -550,6 +556,11 @@ async def test_psycopg_async_avg(self):
 
         await engine.dispose()
 
+
+class TestSqlalchemyAsync2:
+    def setup_method(self):
+        delete_items()
+
     @pytest.mark.asyncio
     @pytest.mark.skipif(sqlalchemy_version == 1, reason='Requires SQLAlchemy 2+')
     async def test_psycopg_async_vector_array(self):