LightRAG/starter/run_all_tests.sh
2025-12-05 14:31:13 +08:00

370 lines
10 KiB
Bash

#!/bin/bash
# ============================================================================
# LightRAG Multi-Tenant Testing Script
#
# This script runs all three testing scenarios sequentially:
# 1. Backward Compatibility Mode (MULTITENANT_MODE=off)
# 2. Single-Tenant Multi-KB Mode (MULTITENANT_MODE=on)
# 3. Full Multi-Tenant Demo Mode (MULTITENANT_MODE=demo)
#
# Usage: ./run_all_tests.sh
# ============================================================================
set -e
# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color
# Configuration
PROJECT_NAME="lightrag-multitenant"
DOCKER_COMPOSE_FILE="docker-compose.yml"
LOG_DIR="test_results"
TIMESTAMP=$(date +%Y%m%d_%H%M%S)
# Create results directory
mkdir -p "$LOG_DIR"
# ============================================================================
# Helper Functions
# ============================================================================
log_header() {
echo -e "${BLUE}╔════════════════════════════════════════════════════════════════════╗${NC}"
echo -e "${BLUE}$1${NC}"
echo -e "${BLUE}╚════════════════════════════════════════════════════════════════════╝${NC}"
}
log_success() {
echo -e "${GREEN}$1${NC}"
}
log_error() {
echo -e "${RED}$1${NC}"
}
log_warning() {
echo -e "${YELLOW}$1${NC}"
}
log_info() {
echo -e "${BLUE}$1${NC}"
}
# ============================================================================
# Pre-flight Checks
# ============================================================================
preflight_checks() {
log_header "PRE-FLIGHT CHECKS"
# Check Docker
if ! command -v docker &> /dev/null; then
log_error "Docker is not installed"
exit 1
fi
log_success "Docker is installed"
# Check Docker Compose
if ! command -v docker compose &> /dev/null; then
log_error "Docker Compose is not installed"
exit 1
fi
log_success "Docker Compose is installed"
# Check pytest
if ! command -v pytest &> /dev/null; then
log_error "pytest is not installed"
exit 1
fi
log_success "pytest is installed"
# Check psql for database connections
if ! command -v psql &> /dev/null; then
log_warning "psql is not installed (optional, for manual DB inspection)"
fi
echo ""
}
# ============================================================================
# Environment Setup
# ============================================================================
setup_environment() {
local mode=$1
log_header "SETTING UP ENVIRONMENT FOR MODE: $mode"
# Create .env file
cp env.example .env
# Add mode-specific configuration
echo "MULTITENANT_MODE=$mode" >> .env
if [ "$mode" = "on" ]; then
echo "DEFAULT_TENANT=tenant-1" >> .env
echo "CREATE_DEFAULT_KB=kb-default,kb-secondary,kb-experimental" >> .env
fi
log_success "Environment configured for mode: $mode"
echo ""
}
# ============================================================================
# Docker Service Management
# ============================================================================
start_services() {
local mode=$1
log_header "STARTING SERVICES (Mode: $mode)"
log_info "Starting Docker Compose services..."
docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME up -d
log_info "Waiting for services to be healthy..."
sleep 15
# Check if services are running
if ! docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME ps | grep -q "healthy"; then
log_warning "Some services may not be healthy yet, but proceeding..."
fi
log_success "Services started"
echo ""
}
stop_services() {
log_info "Stopping services..."
docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME down
log_success "Services stopped"
}
init_database() {
local mode=$1
log_header "INITIALIZING DATABASE (Mode: $mode)"
log_info "Creating database schema..."
docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME exec -T postgres \
psql -U lightrag -d postgres \
-c "CREATE DATABASE lightrag_multitenant;" 2>/dev/null || true
log_info "Applying initialization script..."
docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME exec -T postgres \
psql -U lightrag -d lightrag_multitenant \
-f /docker-entrypoint-initdb.d/01-init.sql
log_success "Database initialized"
echo ""
}
# ============================================================================
# Test Execution
# ============================================================================
run_tests() {
local mode=$1
local log_file="$LOG_DIR/test_${mode}_${TIMESTAMP}.log"
log_header "RUNNING TESTS (Mode: $mode)"
log_info "Executing pytest..."
export MULTITENANT_MODE=$mode
# Select test files based on mode
local test_files=""
if [ "$mode" = "off" ]; then
test_files="tests/test_backward_compatibility.py"
elif [ "$mode" = "on" ]; then
test_files="tests/test_multi_tenant_backends.py::TestTenantIsolation"
elif [ "$mode" = "demo" ]; then
test_files="tests/test_multi_tenant_backends.py tests/test_tenant_security.py"
fi
if pytest $test_files -v --tb=short 2>&1 | tee "$log_file"; then
log_success "All tests passed for mode: $mode"
echo "$mode: PASSED" >> "$LOG_DIR/summary_${TIMESTAMP}.txt"
else
log_error "Tests failed for mode: $mode"
echo "$mode: FAILED" >> "$LOG_DIR/summary_${TIMESTAMP}.txt"
return 1
fi
echo ""
}
# ============================================================================
# Test Scenarios
# ============================================================================
test_scenario_1() {
log_header "SCENARIO 1: BACKWARD COMPATIBILITY MODE (MULTITENANT_MODE=off)"
setup_environment "off"
start_services "off"
init_database "off"
if ! run_tests "off"; then
log_error "Scenario 1 tests failed"
stop_services
return 1
fi
stop_services
log_success "Scenario 1 completed successfully"
echo ""
}
test_scenario_2() {
log_header "SCENARIO 2: SINGLE-TENANT MULTI-KB MODE (MULTITENANT_MODE=on)"
setup_environment "on"
start_services "on"
init_database "on"
if ! run_tests "on"; then
log_error "Scenario 2 tests failed"
stop_services
return 1
fi
stop_services
log_success "Scenario 2 completed successfully"
echo ""
}
test_scenario_3() {
log_header "SCENARIO 3: FULL MULTI-TENANT DEMO MODE (MULTITENANT_MODE=demo)"
setup_environment "demo"
start_services "demo"
init_database "demo"
if ! run_tests "demo"; then
log_error "Scenario 3 tests failed"
stop_services
return 1
fi
stop_services
log_success "Scenario 3 completed successfully"
echo ""
}
# ============================================================================
# Health Checks
# ============================================================================
health_check() {
local mode=$1
log_header "HEALTH CHECK (Mode: $mode)"
log_info "Checking API health..."
if curl -s http://localhost:8000/health > /dev/null; then
log_success "API is healthy"
else
log_warning "API health check failed"
fi
log_info "Checking database..."
if docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME exec -T postgres \
pg_isready -U lightrag -d lightrag_multitenant > /dev/null; then
log_success "Database is ready"
else
log_warning "Database not ready"
fi
echo ""
}
# ============================================================================
# Cleanup
# ============================================================================
cleanup() {
log_header "CLEANUP"
log_info "Removing Docker containers..."
docker compose -f $DOCKER_COMPOSE_FILE -p $PROJECT_NAME down -v 2>/dev/null || true
log_info "Removing temporary .env file..."
rm -f .env
log_success "Cleanup completed"
echo ""
}
# ============================================================================
# Main Execution
# ============================================================================
main() {
log_header "LIGHTRAG MULTI-TENANT TESTING SUITE"
echo ""
echo "Testing Strategy:"
echo " • Scenario 1: Backward Compatibility (no multi-tenant)"
echo " • Scenario 2: Single-Tenant Multi-KB"
echo " • Scenario 3: Full Multi-Tenant"
echo ""
echo "Results will be saved to: $LOG_DIR/"
echo ""
# Create summary file
> "$LOG_DIR/summary_${TIMESTAMP}.txt"
echo "Test Summary - $TIMESTAMP" > "$LOG_DIR/summary_${TIMESTAMP}.txt"
echo "==========================================" >> "$LOG_DIR/summary_${TIMESTAMP}.txt"
echo "" >> "$LOG_DIR/summary_${TIMESTAMP}.txt"
# Run preflight checks
preflight_checks
local failed_tests=0
# Run all scenarios
if ! test_scenario_1; then
((failed_tests++))
log_error "Scenario 1 FAILED"
fi
if ! test_scenario_2; then
((failed_tests++))
log_error "Scenario 2 FAILED"
fi
if ! test_scenario_3; then
((failed_tests++))
log_error "Scenario 3 FAILED"
fi
# Cleanup
cleanup
# Final Report
log_header "TEST SUMMARY"
if [ $failed_tests -eq 0 ]; then
log_success "ALL TESTS PASSED"
echo ""
echo -e "${GREEN}✓ Backward Compatibility: PASSED${NC}"
echo -e "${GREEN}✓ Single-Tenant Multi-KB: PASSED${NC}"
echo -e "${GREEN}✓ Full Multi-Tenant: PASSED${NC}"
echo ""
cat "$LOG_DIR/summary_${TIMESTAMP}.txt"
exit 0
else
log_error "$failed_tests SCENARIO(S) FAILED"
echo ""
cat "$LOG_DIR/summary_${TIMESTAMP}.txt"
exit 1
fi
}
# Handle signals
trap cleanup EXIT
# Run main function
main "$@"