Improved typechecking for tests [skip ci]

Author: ankane

tests/test_asyncpg.py

@@ -64,7 +64,7 @@ async def test_bit(self):
 
         await register_vector(conn)
 
-        embedding = asyncpg.BitString('101')
+        embedding = asyncpg.BitString('101')  # type: ignore
         await conn.execute("INSERT INTO asyncpg_items (embedding) VALUES ($1), (NULL)", embedding)
 
         res = await conn.fetch("SELECT * FROM asyncpg_items ORDER BY id")

tests/test_sqlalchemy.py

@@ -190,7 +190,7 @@ def test_vector(self, engine):
         with Session(engine) as session:
             session.add(Item(id=1, embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert np.array_equal(item.embedding, [1, 2, 3])
 
     def test_vector_l2_distance(self, engine):
@@ -245,7 +245,7 @@ def test_halfvec(self, engine):
         with Session(engine) as session:
             session.add(Item(id=1, half_embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.half_embedding == HalfVector([1, 2, 3])
 
     def test_halfvec_l2_distance(self, engine):
@@ -300,7 +300,7 @@ def test_bit(self, engine):
         with Session(engine) as session:
             session.add(Item(id=1, binary_embedding='101'))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.binary_embedding == '101'
 
     def test_bit_hamming_distance(self, engine):
@@ -337,7 +337,7 @@ def test_sparsevec(self, engine):
         with Session(engine) as session:
             session.add(Item(id=1, sparse_embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.sparse_embedding == SparseVector([1, 2, 3])
 
     def test_sparsevec_l2_distance(self, engine):
@@ -560,7 +560,7 @@ def test_vector_array(self, engine):
             session.commit()
 
             # this fails if the driver does not cast arrays
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert np.array_equal(item.embeddings[0], [1, 2, 3])
             assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
@@ -570,7 +570,7 @@ def test_halfvec_array(self, engine):
             session.commit()
 
             # this fails if the driver does not cast arrays
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.half_embeddings == [HalfVector([1, 2, 3]), HalfVector([4, 5, 6])]
 
 
@@ -587,7 +587,7 @@ async def test_vector(self, engine):
             async with session.begin():
                 embedding = np.array([1, 2, 3])
                 session.add(Item(id=1, embedding=embedding))
-                item = await session.get(Item, 1)
+                item = await session.get_one(Item, 1)
                 assert np.array_equal(item.embedding, embedding)
 
         await engine.dispose()
@@ -600,7 +600,7 @@ async def test_halfvec(self, engine):
             async with session.begin():
                 embedding = [1, 2, 3]
                 session.add(Item(id=1, half_embedding=embedding))
-                item = await session.get(Item, 1)
+                item = await session.get_one(Item, 1)
                 assert item.half_embedding == HalfVector(embedding)
 
         await engine.dispose()
@@ -613,12 +613,12 @@ async def test_bit(self, engine):
             async with session.begin():
                 embedding = asyncpg.BitString('101') if engine == asyncpg_engine else '101'
                 session.add(Item(id=1, binary_embedding=embedding))
-                item = await session.get(Item, 1)
+                item = await session.get_one(Item, 1)
                 assert item.binary_embedding == embedding
 
                 if engine == asyncpg_engine:
                     session.add(Item(id=2, binary_embedding='101'))
-                    item = await session.get(Item, 2)
+                    item = await session.get_one(Item, 2)
                     assert item.binary_embedding == embedding
 
         await engine.dispose()
@@ -631,7 +631,7 @@ async def test_sparsevec(self, engine):
             async with session.begin():
                 embedding = [1, 2, 3]
                 session.add(Item(id=1, sparse_embedding=embedding))
-                item = await session.get(Item, 1)
+                item = await session.get_one(Item, 1)
                 assert item.sparse_embedding == SparseVector(embedding)
 
         await engine.dispose()
@@ -662,12 +662,12 @@ async def test_vector_array(self, engine):
         async with async_session() as session:
             async with session.begin():
                 session.add(Item(id=1, embeddings=[Vector([1, 2, 3]), Vector([4, 5, 6])]))
-                item = await session.get(Item, 1)
+                item = await session.get_one(Item, 1)
                 assert np.array_equal(item.embeddings[0], [1, 2, 3])
                 assert np.array_equal(item.embeddings[1], [4, 5, 6])
 
                 session.add(Item(id=2, embeddings=[np.array([1, 2, 3]), np.array([4, 5, 6])]))
-                item = await session.get(Item, 2)
+                item = await session.get_one(Item, 2)
                 assert np.array_equal(item.embeddings[0], [1, 2, 3])
                 assert np.array_equal(item.embeddings[1], [4, 5, 6])
 

tests/test_sqlmodel.py

@@ -75,7 +75,7 @@ def test_vector(self):
         with Session(engine) as session:
             session.add(Item(id=1, embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert np.array_equal(item.embedding, np.array([1, 2, 3]))
 
     def test_vector_l2_distance(self):
@@ -106,7 +106,7 @@ def test_halfvec(self):
         with Session(engine) as session:
             session.add(Item(id=1, half_embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.half_embedding == HalfVector([1, 2, 3])
 
     def test_halfvec_l2_distance(self):
@@ -137,7 +137,7 @@ def test_bit(self):
         with Session(engine) as session:
             session.add(Item(id=1, binary_embedding='101'))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.binary_embedding == '101'
 
     def test_bit_hamming_distance(self):
@@ -156,7 +156,7 @@ def test_sparsevec(self):
         with Session(engine) as session:
             session.add(Item(id=1, sparse_embedding=[1, 2, 3]))
             session.commit()
-            item = session.get(Item, 1)
+            item = session.get_one(Item, 1)
             assert item.sparse_embedding == SparseVector([1, 2, 3])
 
     def test_sparsevec_l2_distance(self):