tests/unit/test_vector_db.py

"""
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()