Newer
Older
cortex-hub / ai-hub / run_integration_tests.sh
#!/bin/bash

# A script to automate running tests locally.
# It starts the FastAPI server, runs the specified tests, and then shuts down the server.
#
# To run without starting the server, use the --skip-server flag:
# ./run_tests.sh --skip-server

# --- Configuration ---
# You can define aliases for your test file paths here.
TEST_SUITES=(
  "All tests"
  "integration_tests/test_sessions_api.py"
  "integration_tests/test_documents_api.py"
  "integration_tests/test_misc_api.py"
)
TEST_PATHS=(
  "integration_tests/"
  "integration_tests/test_sessions_api.py"
  "integration_tests/test_documents_api.py"
  "integration_tests/test_misc_api.py"
)

export DB_MODE=sqlite
export LOCAL_DB_PATH="data/integration_test_ai_hub.db"
export FAISS_INDEX_PATH="data/integration_test_faiss_index.bin"
export LOG_LEVEL="DEBUG"

# Check for the --skip-server flag as an inline parameter.
# The user can pass this as the first argument to the script.
SKIP_SERVER=false
if [ "$1" == "--skip-server" ]; then
    SKIP_SERVER=true
    shift # Remove the --skip-server argument from the list
fi

# --- User Interaction ---
echo "--- AI Hub Test Runner ---"
echo "Select a test suite to run:"

# Display the menu options
for i in "${!TEST_SUITES[@]}"; do
  echo "$((i+1)). ${TEST_SUITES[$i]}"
done

# Prompt for user input
read -p "Enter the number of your choice (1-${#TEST_SUITES[@]}): " choice

# Validate input and set the TEST_PATH
if [[ "$choice" =~ ^[0-9]+$ ]] && [ "$choice" -ge 1 ] && [ "$choice" -le "${#TEST_SUITES[@]}" ]; then
  TEST_PATH=${TEST_PATHS[$((choice-1))]}
  echo "You have selected: ${TEST_SUITES[$((choice-1))]}"
else
  echo "Invalid choice. Running all tests by default."
  TEST_PATH=${TEST_PATHS[0]}
fi

# --- Conditional Server Startup and Cleanup ---
if [ "$SKIP_SERVER" = false ]; then
    # --- Pre-test Cleanup ---
    # Check for and remove old test files to ensure a clean test environment.
    echo "--- Checking for and removing old test files ---"
    if [ -f "$LOCAL_DB_PATH" ]; then
        echo "Removing old database file: $LOCAL_DB_PATH"
        rm "$LOCAL_DB_PATH"
    fi
    if [ -f "$FAISS_INDEX_PATH" ]; then
        echo "Removing old FAISS index file: $FAISS_INDEX_PATH"
        rm "$FAISS_INDEX_PATH"
    fi
    echo "Cleanup complete."

    echo "--- Starting AI Hub Server for Tests ---"

    # Start the uvicorn server in the background
    # We bind it to 127.0.0.1 to ensure it's not accessible from outside the local machine.
    uvicorn app.main:app --host 127.0.0.1 --port 8001 --reload &

    # Get the Process ID (PID) of the background server
    SERVER_PID=$!

    # Define a cleanup function to be called on exit
    cleanup() {
        echo ""
        echo "--- Shutting Down Server (PID: $SERVER_PID) ---"
        kill $SERVER_PID
    }

    # Register the cleanup function to run when the script exits
    # This ensures the server is stopped even if tests fail or the script is interrupted (e.g., with Ctrl+C).
    trap cleanup EXIT

    echo "Server started with PID: $SERVER_PID. Waiting for it to initialize..."

    # Wait a few seconds to ensure the server is fully up and running
    sleep 5
else
    echo "--- Skipping server startup. Running tests against an existing server. ---"
fi

echo "--- Running tests in: $TEST_PATH ---"

# Execute the Python tests using pytest on the specified path
# The '-s' flag shows print statements from the tests.
pytest -s "$TEST_PATH"

# Capture the exit code of the test script
TEST_EXIT_CODE=$?

# The 'trap' will automatically call the cleanup function if it was set.
# Exit with the same code as the test script (0 for success, non-zero for failure).
exit $TEST_EXIT_CODE