forked from innovacion/searchbox
56 lines
1.4 KiB
Python
56 lines
1.4 KiB
Python
from uuid import uuid4
|
|
import pytest
|
|
from searchbox.client import Client
|
|
from searchbox.engine.qdrant_engine import QdrantEngine
|
|
from searchbox.models import Chunk
|
|
|
|
|
|
@pytest.fixture(scope="module")
|
|
def qdrant_client():
|
|
client = Client("qdrant", "test_index", url=":memory:")
|
|
return client
|
|
|
|
|
|
def test_create_qdrant_client(qdrant_client: Client):
|
|
assert isinstance(qdrant_client, Client)
|
|
assert isinstance(qdrant_client.engine, QdrantEngine)
|
|
|
|
|
|
async def test_create_qdrant_index(qdrant_client: Client):
|
|
result = await qdrant_client.create_index(3)
|
|
|
|
assert result is True
|
|
|
|
|
|
async def test_upload_chunk(qdrant_client: Client):
|
|
chunk = Chunk.model_validate(
|
|
{
|
|
"id": uuid4(),
|
|
"vector": [0.0, 0.1, 0.3],
|
|
"payload": {
|
|
"page_content": "This is a test page content.",
|
|
"filename": "test.txt",
|
|
"page": 1,
|
|
},
|
|
}
|
|
)
|
|
|
|
result = await qdrant_client.upload_chunk(chunk)
|
|
|
|
assert result is True
|
|
|
|
|
|
async def test_semantic_search(qdrant_client: Client):
|
|
result = await qdrant_client.semantic_search([0.0, 0.1, 0.3])
|
|
|
|
assert len(result) == 1
|
|
|
|
first_result = result[0]
|
|
assert first_result.chunk_id is not None
|
|
assert first_result.score > 0.9
|
|
assert first_result.payload == {
|
|
"page_content": "This is a test page content.",
|
|
"filename": "test.txt",
|
|
"page": 1,
|
|
}
|