40 lines
1.2 KiB
Python
40 lines
1.2 KiB
Python
import pytest
|
|
from pathlib import Path
|
|
|
|
|
|
def test_chroma_client_creates_persist_directory(tmp_path):
|
|
import os
|
|
os.environ["CHROMA_DB_PATH"] = str(tmp_path / "test_chroma")
|
|
|
|
from app.core.database import get_chroma_client
|
|
|
|
client = get_chroma_client()
|
|
assert client is not None
|
|
assert (tmp_path / "test_chroma").exists()
|
|
|
|
|
|
def test_chroma_client_creates_new_collection(tmp_path):
|
|
import os
|
|
os.environ["CHROMA_DB_PATH"] = str(tmp_path / "test_chroma")
|
|
|
|
from app.core.database import get_chroma_client, get_or_create_collection
|
|
|
|
client = get_chroma_client()
|
|
collection = get_or_create_collection(client, "test_docs")
|
|
assert collection.name == "test_docs"
|
|
assert collection.count() == 0
|
|
|
|
|
|
def test_chroma_client_returns_existing_collection(tmp_path):
|
|
import os
|
|
os.environ["CHROMA_DB_PATH"] = str(tmp_path / "test_chroma")
|
|
|
|
from app.core.database import get_chroma_client, get_or_create_collection
|
|
|
|
client = get_chroma_client()
|
|
collection1 = get_or_create_collection(client, "test_docs")
|
|
collection1.add(documents=["test"], ids=["1"])
|
|
|
|
collection2 = get_or_create_collection(client, "test_docs")
|
|
assert collection2.count() == 1
|