File size: 1,393 Bytes
939c0c0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 | """
Unit tests for text chunking, embedding, reranking, and vector DB functionality.
"""
from __future__ import annotations
import pytest
import numpy as np
from app.services.document_service import chunk_text, _clean_text
from services.reranker import RankedChunk, reranker_service
def test_clean_and_chunk_text():
text = "Hello world! " * 100
cleaned = _clean_text(text)
assert len(cleaned) > 0
chunks = chunk_text(text, chunk_size=200, overlap=20)
assert len(chunks) > 1
assert all(len(c) <= 220 for c in chunks)
@pytest.mark.asyncio
async def test_reranker_service():
query = "What is the vacation policy?"
chunks = [
RankedChunk(
text="Employees get 20 days of paid vacation annually.",
score=0.5,
document_id="doc1",
document_name="hr.txt",
page=1,
doc_type="hr"
),
RankedChunk(
text="The server database connection string is configured in env.",
score=0.8,
document_id="doc2",
document_name="tech.txt",
page=1,
doc_type="technical"
)
]
# Rerank
reranked = await reranker_service.rerank(query, chunks, top_k=2)
assert len(reranked) == 2
# The vacation-related chunk should score higher after reranking
assert "vacation" in reranked[0].text.lower()
|