import tempfile

import pytest

from companion.rag.vector_store import VectorStore


def test_vector_store_upsert_and_search():
    with tempfile.TemporaryDirectory() as tmp:
        store = VectorStore(uri=tmp, dimensions=4)
        store.upsert(
            ids=["a", "b"],
            texts=["hello world", "goodbye world"],
            embeddings=[[1.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0]],
            metadatas=[
                {"source_file": "a.md", "source_directory": "docs"},
                {"source_file": "b.md", "source_directory": "docs"},
            ],
        )
        results = store.search(query_vector=[1.0, 0.0, 0.0, 0.0], top_k=1)
        assert len(results) == 1
        assert results[0]["source_file"] == "a.md"