283 lines
8.3 KiB
Python
283 lines
8.3 KiB
Python
"""
|
|
Pytest configuration and shared fixtures for Code-Forge tests.
|
|
"""
|
|
|
|
import os
|
|
import pytest
|
|
import psycopg2
|
|
from typing import Generator, Dict, Any
|
|
from unittest.mock import Mock, MagicMock
|
|
import numpy as np
|
|
|
|
# Set test environment variables
|
|
os.environ["POSTGRES_HOST"] = os.getenv("TEST_POSTGRES_HOST", "localhost")
|
|
os.environ["POSTGRES_PORT"] = os.getenv("TEST_POSTGRES_PORT", "5432")
|
|
os.environ["POSTGRES_DB"] = os.getenv("TEST_POSTGRES_DB", "codeforge_test")
|
|
os.environ["POSTGRES_USER"] = os.getenv("TEST_POSTGRES_USER", os.getenv("USER", "postgres"))
|
|
os.environ["OLLAMA_HOST"] = os.getenv("TEST_OLLAMA_HOST", "http://localhost:11434")
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def test_db_config() -> Dict[str, str]:
|
|
"""Database configuration for tests."""
|
|
return {
|
|
"host": os.getenv("POSTGRES_HOST"),
|
|
"port": int(os.getenv("POSTGRES_PORT")),
|
|
"database": os.getenv("POSTGRES_DB"),
|
|
"user": os.getenv("POSTGRES_USER"),
|
|
}
|
|
|
|
|
|
@pytest.fixture(scope="session")
|
|
def db_connection(test_db_config):
|
|
"""
|
|
Create a test database connection.
|
|
This is a session-scoped fixture that creates the test database once.
|
|
"""
|
|
# Connect to default postgres database to create test database
|
|
conn = psycopg2.connect(
|
|
host=test_db_config["host"],
|
|
port=test_db_config["port"],
|
|
database="postgres",
|
|
user=test_db_config["user"],
|
|
)
|
|
conn.autocommit = True
|
|
|
|
with conn.cursor() as cur:
|
|
# Drop and recreate test database
|
|
cur.execute(f"DROP DATABASE IF EXISTS {test_db_config['database']}")
|
|
cur.execute(f"CREATE DATABASE {test_db_config['database']}")
|
|
|
|
conn.close()
|
|
|
|
# Connect to test database
|
|
test_conn = psycopg2.connect(**test_db_config)
|
|
|
|
# Create pgvector extension
|
|
with test_conn.cursor() as cur:
|
|
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
|
|
test_conn.commit()
|
|
|
|
# Load schema
|
|
with open("app/core/schema_pgvector.sql", "r") as f:
|
|
schema_sql = f.read()
|
|
|
|
with test_conn.cursor() as cur:
|
|
cur.execute(schema_sql)
|
|
test_conn.commit()
|
|
|
|
yield test_conn
|
|
|
|
# Cleanup
|
|
test_conn.close()
|
|
|
|
# Drop test database
|
|
conn = psycopg2.connect(
|
|
host=test_db_config["host"],
|
|
port=test_db_config["port"],
|
|
database="postgres",
|
|
user=test_db_config["user"],
|
|
)
|
|
conn.autocommit = True
|
|
|
|
with conn.cursor() as cur:
|
|
cur.execute(f"DROP DATABASE IF EXISTS {test_db_config['database']}")
|
|
|
|
conn.close()
|
|
|
|
|
|
@pytest.fixture
|
|
def clean_db(db_connection):
|
|
"""
|
|
Clean database before each test.
|
|
This is a function-scoped fixture that cleans tables before each test.
|
|
Works with NEW PGVectorStore schema (data_{kb_name} tables).
|
|
"""
|
|
with db_connection.cursor() as cur:
|
|
# Get all KB table names
|
|
cur.execute("SELECT table_name FROM knowledge_bases WHERE table_name IS NOT NULL")
|
|
kb_tables = cur.fetchall()
|
|
|
|
# Truncate each KB data table
|
|
for (table_name,) in kb_tables:
|
|
try:
|
|
cur.execute(f"TRUNCATE TABLE {table_name} CASCADE")
|
|
except Exception:
|
|
# Table might not exist, skip it
|
|
pass
|
|
|
|
# Truncate knowledge_bases table
|
|
cur.execute("TRUNCATE TABLE knowledge_bases CASCADE")
|
|
db_connection.commit()
|
|
|
|
yield db_connection
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_kb(clean_db):
|
|
"""Create a sample knowledge base for testing with NEW PGVectorStore schema."""
|
|
import json
|
|
|
|
table_name = "data_test_kb"
|
|
embed_dim = 768
|
|
|
|
with clean_db.cursor() as cur:
|
|
# Insert KB metadata with NEW schema columns
|
|
# Note: metadata must be JSON string, not dict
|
|
cur.execute(
|
|
"""
|
|
INSERT INTO knowledge_bases (name, kb_type, description, metadata, table_name, embed_dim)
|
|
VALUES (%s, %s, %s, %s, %s, %s)
|
|
RETURNING id
|
|
""",
|
|
("test_kb", "GENERIC", "Test knowledge base", json.dumps({}), table_name, embed_dim)
|
|
)
|
|
kb_id = cur.fetchone()[0]
|
|
|
|
# Create the actual data table using helper function
|
|
cur.execute("SELECT create_kb_table(%s, %s, %s)", ("test_kb", table_name, embed_dim))
|
|
|
|
clean_db.commit()
|
|
|
|
return {
|
|
"id": kb_id,
|
|
"name": "test_kb",
|
|
"kb_type": "GENERIC",
|
|
"table_name": table_name,
|
|
"embed_dim": embed_dim
|
|
}
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_embedding() -> list:
|
|
"""Generate a sample 768-dimensional embedding."""
|
|
np.random.seed(42) # For reproducibility
|
|
return np.random.randn(768).tolist()
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_documents(sample_kb, sample_embedding, clean_db):
|
|
"""Create sample documents with embeddings in NEW PGVectorStore schema."""
|
|
import json
|
|
|
|
documents = []
|
|
table_name = sample_kb["table_name"] # Should be 'data_test_kb'
|
|
|
|
for i in range(5):
|
|
# Generate slightly different embeddings
|
|
np.random.seed(42 + i)
|
|
embedding = np.random.randn(768).tolist()
|
|
node_id = f"node_doc_{i}"
|
|
text_content = f"This is test document {i} with some content about testing."
|
|
metadata = {"filename": f"test_{i}.txt", "chunk_index": i}
|
|
|
|
with clean_db.cursor() as cur:
|
|
# Insert into data_test_kb table with NEW schema columns
|
|
# Note: metadata_ is JSON type, needs JSON string
|
|
cur.execute(
|
|
f"""
|
|
INSERT INTO {table_name} (text, metadata_, node_id, embedding)
|
|
VALUES (%s, %s, %s, %s)
|
|
RETURNING id
|
|
""",
|
|
(text_content, json.dumps(metadata), node_id, embedding)
|
|
)
|
|
doc_id = cur.fetchone()[0]
|
|
|
|
documents.append({
|
|
"id": doc_id,
|
|
"node_id": node_id,
|
|
"text": text_content,
|
|
"metadata": metadata,
|
|
"embedding": embedding
|
|
})
|
|
|
|
clean_db.commit()
|
|
|
|
return documents
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ollama_embedding():
|
|
"""Mock Ollama embedding model."""
|
|
mock = MagicMock()
|
|
mock.get_text_embedding.return_value = np.random.randn(768).tolist()
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def mock_ollama_llm():
|
|
"""Mock Ollama LLM."""
|
|
mock = MagicMock()
|
|
mock_response = MagicMock()
|
|
mock_response.message.content = "This is a test response from the LLM."
|
|
mock.chat.return_value = mock_response
|
|
return mock
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_text_file(tmp_path):
|
|
"""Create a sample text file for testing."""
|
|
file_path = tmp_path / "test_document.txt"
|
|
file_path.write_text("This is a test document.\n\nIt has multiple paragraphs.\n\nFor testing purposes.")
|
|
return file_path
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_pdf_file(tmp_path):
|
|
"""Create a sample PDF file for testing (requires PyMuPDF)."""
|
|
try:
|
|
import fitz # PyMuPDF
|
|
|
|
file_path = tmp_path / "test_document.pdf"
|
|
doc = fitz.open()
|
|
page = doc.new_page()
|
|
page.insert_text((72, 72), "This is a test PDF document.\n\nIt has multiple paragraphs.\n\nFor testing purposes.")
|
|
doc.save(str(file_path))
|
|
doc.close()
|
|
|
|
return file_path
|
|
except ImportError:
|
|
pytest.skip("PyMuPDF not installed")
|
|
|
|
|
|
@pytest.fixture
|
|
def sample_markdown_file(tmp_path):
|
|
"""Create a sample Markdown file for testing."""
|
|
file_path = tmp_path / "test_document.md"
|
|
content = """# Test Document
|
|
|
|
This is a test document.
|
|
|
|
## Section 1
|
|
|
|
Some content here.
|
|
|
|
## Section 2
|
|
|
|
More content here.
|
|
"""
|
|
file_path.write_text(content)
|
|
return file_path
|
|
|
|
|
|
@pytest.fixture
|
|
def api_client():
|
|
"""Create a test client for FastAPI."""
|
|
from fastapi.testclient import TestClient
|
|
from mcp_server.server import app
|
|
|
|
return TestClient(app)
|
|
|
|
|
|
# Markers for test categorization
|
|
def pytest_configure(config):
|
|
"""Configure custom markers."""
|
|
config.addinivalue_line("markers", "unit: Unit tests (fast, no external dependencies)")
|
|
config.addinivalue_line("markers", "integration: Integration tests (require database, Ollama)")
|
|
config.addinivalue_line("markers", "slow: Slow tests (may take several seconds)")
|
|
config.addinivalue_line("markers", "embeddings: Tests involving embedding generation")
|
|
config.addinivalue_line("markers", "vector: Tests involving vector operations")
|
|
config.addinivalue_line("markers", "rag: Tests involving RAG engine")
|
|
config.addinivalue_line("markers", "api: Tests involving API endpoints")
|
|
|