From 00570fde316e3683023b736257898b8e5ba9788a Mon Sep 17 00:00:00 2001 From: Ben Browning Date: Wed, 5 Mar 2025 13:20:13 -0500 Subject: [PATCH] chore: Get sqlite_vec and vector_store unit tests passing (#1413) --- tests/unit/providers/vector_io/test_sqlite_vec.py | 3 ++- tests/unit/rag/test_vector_store.py | 8 +++++--- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/tests/unit/providers/vector_io/test_sqlite_vec.py b/tests/unit/providers/vector_io/test_sqlite_vec.py index e1d87de24..eb5660a85 100644 --- a/tests/unit/providers/vector_io/test_sqlite_vec.py +++ b/tests/unit/providers/vector_io/test_sqlite_vec.py @@ -9,6 +9,7 @@ import sqlite3 import numpy as np import pytest +import pytest_asyncio import sqlite_vec from llama_stack.apis.vector_io import Chunk, QueryChunksResponse @@ -48,7 +49,7 @@ def sqlite_connection(loop): conn.close() -@pytest.fixture(scope="session", autouse=True) +@pytest_asyncio.fixture(scope="session", autouse=True) async def sqlite_vec_index(sqlite_connection): return await SQLiteVecIndex.create(dimension=EMBEDDING_DIMENSION, connection=sqlite_connection, bank_id="test_bank") diff --git a/tests/unit/rag/test_vector_store.py b/tests/unit/rag/test_vector_store.py index e0d340657..3decc431e 100644 --- a/tests/unit/rag/test_vector_store.py +++ b/tests/unit/rag/test_vector_store.py @@ -15,6 +15,8 @@ from llama_stack.apis.tools import RAGDocument from llama_stack.providers.utils.memory.vector_store import URL, content_from_doc DUMMY_PDF_PATH = Path(os.path.abspath(__file__)).parent / "fixtures" / "dummy.pdf" +# Depending on the machine, this can get parsed a couple of ways +DUMMY_PDF_TEXT_CHOICES = ["Dummy PDF file", "Dumm y PDF file"] def read_file(file_path: str) -> bytes: @@ -45,7 +47,7 @@ class TestVectorStore: metadata={}, ) content = await content_from_doc(doc) - assert content == "Dumm y PDF file" + assert content in DUMMY_PDF_TEXT_CHOICES @pytest.mark.asyncio async def test_downloads_pdf_and_returns_content(self): @@ -58,7 +60,7 @@ class TestVectorStore: metadata={}, ) content = await content_from_doc(doc) - assert content == "Dumm y PDF file" + assert content in DUMMY_PDF_TEXT_CHOICES @pytest.mark.asyncio async def test_downloads_pdf_and_returns_content_with_url_object(self): @@ -73,4 +75,4 @@ class TestVectorStore: metadata={}, ) content = await content_from_doc(doc) - assert content == "Dumm y PDF file" + assert content in DUMMY_PDF_TEXT_CHOICES