370 lines
10 KiB
Bash
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 "$@"
|