44 lines
1.1 KiB
Python
44 lines
1.1 KiB
Python
import os
|
|
import tempfile
|
|
import shutil
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import chromadb
|
|
from chromadb.config import Settings
|
|
|
|
from src.ingest import ingest_faq
|
|
|
|
def test_ingest_faq(tmp_path):
|
|
# Create a temporary FAQ file
|
|
faq_content = """Q: What is Python?
|
|
A: Python is a programming language.
|
|
|
|
Q: What is ChromaDB?
|
|
A: ChromaDB is a vector database."""
|
|
faq_file = tmp_path / "faq.txt"
|
|
faq_file.write_text(faq_content, encoding="utf-8")
|
|
|
|
# Initialize a temporary ChromaDB client
|
|
db_dir = tmp_path / "chromadb"
|
|
client = chromadb.Client(Settings(
|
|
chroma_db_impl="duckdb+parquet",
|
|
persist_directory=str(db_dir)
|
|
))
|
|
|
|
collection_name = "test_collection"
|
|
|
|
# Ingest
|
|
ingest_faq(faq_file, client, collection_name)
|
|
|
|
# Verify collection exists and has documents
|
|
collection = client.get_collection(name=collection_name)
|
|
assert collection.count() == 2
|
|
|
|
# Verify metadata
|
|
docs = collection.get(ids=["0", "1"])
|
|
assert docs["metadatas"][0]["question"] == "What is Python?"
|
|
assert docs["metadatas"][1]["answer"] == "ChromaDB is a vector database."
|
|
|
|
# Clean up
|
|
shutil.rmtree(db_dir) |