67 lines
1.7 KiB
Python
67 lines
1.7 KiB
Python
import pytest
|
|
import shutil
|
|
import tempfile
|
|
import os
|
|
|
|
|
|
@pytest.fixture
|
|
def temp_dir():
|
|
"""Create a temporary directory for test files"""
|
|
temp_dir = tempfile.mkdtemp()
|
|
yield temp_dir
|
|
shutil.rmtree(temp_dir)
|
|
|
|
|
|
@pytest.fixture
|
|
def docs_dir(temp_dir):
|
|
"""Create a temporary documents directory"""
|
|
docs_dir = os.path.join(temp_dir, "documents")
|
|
os.makedirs(docs_dir)
|
|
yield docs_dir
|
|
|
|
|
|
@pytest.fixture
|
|
def db_dir(temp_dir):
|
|
"""Create a temporary vector database directory"""
|
|
db_dir = os.path.join(temp_dir, "vector_db")
|
|
os.makedirs(db_dir)
|
|
yield db_dir
|
|
|
|
|
|
@pytest.fixture
|
|
def tracker_file(temp_dir):
|
|
"""Create a temporary tracker file"""
|
|
tracker_path = os.path.join(temp_dir, "test_tracker.json")
|
|
yield tracker_path
|
|
# Clean up after tests
|
|
if os.path.exists(tracker_path):
|
|
os.remove(tracker_path)
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_docs(docs_dir):
|
|
"""Create sample text documents for testing"""
|
|
# Create a few sample documents
|
|
doc1_path = os.path.join(docs_dir, "doc1.txt")
|
|
doc2_path = os.path.join(docs_dir, "doc2.txt")
|
|
doc3_path = os.path.join(docs_dir, "doc3.txt")
|
|
|
|
with open(doc1_path, "w") as f:
|
|
f.write("This is a sample document about artificial intelligence. " * 10)
|
|
|
|
with open(doc2_path, "w") as f:
|
|
f.write("This document discusses machine learning concepts. " * 10)
|
|
|
|
with open(doc3_path, "w") as f:
|
|
f.write("Natural language processing is a field of AI. " * 10)
|
|
|
|
return [doc1_path, doc2_path, doc3_path]
|
|
|
|
|
|
@pytest.fixture
|
|
def text_splitter():
|
|
"""Create a mock text splitter"""
|
|
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
|
|
|
return RecursiveCharacterTextSplitter(chunk_size=100, chunk_overlap=20)
|