Improved test [skip ci]

ankane · ankane · commit 7cd310b5cf98 · 2025-02-09T13:36:12.000-08:00
diff --git a/tests/test_sqlalchemy.py b/tests/test_sqlalchemy.py
@@ -613,18 +613,17 @@ async def test_sparsevec(self, engine):
 
     @pytest.mark.asyncio
     async def test_avg(self, engine):
-        # TODO do not throw error when types are registered
-        if engine == psycopg_async_type_engine:
-            return
-
         async_session = async_sessionmaker(engine, expire_on_commit=False)
 
         async with async_session() as session:
             async with session.begin():
                 session.add(Item(embedding=[1, 2, 3]))
                 session.add(Item(embedding=[4, 5, 6]))
                 avg = await session.scalars(select(func.avg(Item.embedding)))
-                assert avg.first() == '[2.5,3.5,4.5]'
+                if engine == psycopg_async_type_engine:
+                    assert avg.first().tolist() == [2.5, 3.5, 4.5]
+                else:
+                    assert avg.first() == '[2.5,3.5,4.5]'
 
         await engine.dispose()
 
