Skip to content

Commit 6fde707

Browse files
committed
Fix tests
1 parent 1206c20 commit 6fde707

File tree

3 files changed

+33
-6
lines changed

3 files changed

+33
-6
lines changed

tests/test_async_vectorstore_search.py

Lines changed: 28 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -116,7 +116,34 @@ async def engine(self, db_project, db_region, db_cluster, db_instance, db_name):
116116
await engine.close()
117117

118118
@pytest_asyncio.fixture(scope="class")
119-
async def vs_custom_scann_query_option(self, engine):
119+
async def vs_custom(self, engine):
120+
await engine._ainit_vectorstore_table(
121+
CUSTOM_TABLE,
122+
VECTOR_SIZE,
123+
id_column="myid",
124+
content_column="mycontent",
125+
embedding_column="myembedding",
126+
metadata_columns=[
127+
Column("page", "TEXT"),
128+
Column("source", "TEXT"),
129+
],
130+
store_metadata=False,
131+
)
132+
133+
vs_custom = await AsyncAlloyDBVectorStore.create(
134+
engine,
135+
embedding_service=embeddings_service,
136+
table_name=CUSTOM_TABLE,
137+
id_column="myid",
138+
content_column="mycontent",
139+
embedding_column="myembedding",
140+
index_query_options=HNSWQueryOptions(ef_search=1),
141+
)
142+
await vs_custom.aadd_documents(docs, ids=ids)
143+
yield vs_custom
144+
145+
@pytest_asyncio.fixture(scope="class")
146+
async def vs_custom_scann_query_option(self, engine, vs_custom):
120147
vs_custom_scann_query_option = await AsyncAlloyDBVectorStore.create(
121148
engine,
122149
embedding_service=embeddings_service,

tests/test_vectorstore.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -195,7 +195,7 @@ async def test_aadd_images(self, engine_sync, image_uris):
195195
await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{IMAGE_TABLE}"')
196196

197197
async def test_aadd_images_store_uri_only(self, engine_sync, image_uris):
198-
table_name = IMAGE_TABLE_SYNC + "_store_uri_only"
198+
table_name = IMAGE_TABLE_SYNC + "_uri"
199199
engine_sync.init_vectorstore_table(
200200
table_name,
201201
VECTOR_SIZE,
@@ -255,7 +255,7 @@ async def test_add_images(self, engine_sync, image_uris):
255255
await aexecute(engine_sync, f'DROP TABLE IF EXISTS "{IMAGE_TABLE_SYNC}"')
256256

257257
async def test_add_images_store_uri_only(self, engine_sync, image_uris):
258-
table_name = IMAGE_TABLE_SYNC + "_store_uri_only"
258+
table_name = IMAGE_TABLE_SYNC + "_uri"
259259
engine_sync.init_vectorstore_table(
260260
table_name,
261261
VECTOR_SIZE,

tests/test_vectorstore_embeddings.py

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -183,7 +183,7 @@ async def test_asimilarity_search(self, vs):
183183
results = await vs.asimilarity_search("foo", k=1)
184184
assert len(results) == 1
185185
assert results == [Document(page_content="foo", id=ids[0])]
186-
results = await vs.asimilarity_search("foo", k=1, filter={"mycontent": "bar"})
186+
results = await vs.asimilarity_search("foo", k=1, filter={"content": "bar"})
187187
assert results == [Document(page_content="bar", id=ids[1])]
188188

189189
async def test_asimilarity_search_score(self, vs):
@@ -242,7 +242,7 @@ async def test_amax_marginal_relevance_search(self, vs):
242242
results = await vs.amax_marginal_relevance_search("bar")
243243
assert results[0] == Document(page_content="bar", id=ids[1])
244244
results = await vs.amax_marginal_relevance_search(
245-
"bar", filter={"mycontent": "boo"}
245+
"bar", filter={"content": "boo"}
246246
)
247247
assert results[0] == Document(page_content="boo", id=ids[3])
248248

@@ -343,7 +343,7 @@ def test_similarity_search(self, vs_custom):
343343
assert len(results) == 1
344344
assert results == [Document(page_content="foo", id=ids[0])]
345345
results = vs_custom.similarity_search("foo", k=1, filter={"mycontent": "boo"})
346-
assert results == [Document(page_content="bar", id=ids[1])]
346+
assert results == [Document(page_content="boo", id=ids[3])]
347347

348348
def test_similarity_search_score(self, vs_custom):
349349
results = vs_custom.similarity_search_with_score("foo")

0 commit comments

Comments
 (0)