code_rag/tests/fixtures.py
2025-03-21 10:09:07 -04:00

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)