""" 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")