ai-llm-red-team-handbook/.github/workflows/pit-test.yml

name: PIT - Continuous Integration

on:
  push:
    branches: [main, develop]
    paths:
      - 'tools/prompt_injection_tester/**'
      - '.github/workflows/pit-test.yml'
  pull_request:
    branches: [main, develop]
    paths:
      - 'tools/prompt_injection_tester/**'
  workflow_dispatch:

jobs:
  lint-and-test:
    name: Lint and Test
    runs-on: ubuntu-latest
    strategy:
      matrix:
        python-version: ['3.10', '3.11', '3.12']

    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python ${{ matrix.python-version }}
        uses: actions/setup-python@v5
        with:
          python-version: ${{ matrix.python-version }}
          cache: 'pip'

      - name: Install dependencies
        working-directory: tools/prompt_injection_tester
        run: |
          python -m pip install --upgrade pip
          pip install -e .
          pip install pytest pytest-cov pytest-asyncio black ruff mypy

      - name: Lint with ruff
        working-directory: tools/prompt_injection_tester
        run: |
          ruff check pit/ tests/ --output-format=github
        continue-on-error: true

      - name: Format check with black
        working-directory: tools/prompt_injection_tester
        run: |
          black --check pit/ tests/
        continue-on-error: true

      - name: Type check with mypy
        working-directory: tools/prompt_injection_tester
        run: |
          mypy pit/ --ignore-missing-imports
        continue-on-error: true

      - name: Run integration tests
        working-directory: tools/prompt_injection_tester
        run: |
          pytest tests/integration/ -v --cov=pit --cov-report=xml --cov-report=term

      - name: Run report formatter tests
        working-directory: tools/prompt_injection_tester
        run: |
          python tests/test_reports.py

      - name: Upload coverage reports
        uses: codecov/codecov-action@v4
        with:
          file: tools/prompt_injection_tester/coverage.xml
          flags: unittests
          name: codecov-${{ matrix.python-version }}
        continue-on-error: true

  docker-build:
    name: Docker Build Test
    runs-on: ubuntu-latest

    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Docker Buildx
        uses: docker/setup-buildx-action@v3

      - name: Build Docker image
        working-directory: tools/prompt_injection_tester
        run: |
          docker build -t pit:test .

      - name: Test Docker image
        run: |
          docker run --rm pit:test --version
          docker run --rm pit:test --help

  e2e-test:
    name: End-to-End Test
    runs-on: ubuntu-latest
    needs: [lint-and-test]

    services:
      ollama:
        image: ollama/ollama:latest
        ports:
          - 11434:11434

    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'
          cache: 'pip'

      - name: Install dependencies
        working-directory: tools/prompt_injection_tester
        run: |
          python -m pip install --upgrade pip
          pip install -e .

      - name: Wait for Ollama
        run: |
          timeout 60 bash -c 'until curl -f http://localhost:11434/api/tags; do sleep 2; done'

      - name: Pull Ollama model
        run: |
          docker exec ${{ job.services.ollama.id }} ollama pull llama3.2:1b

      - name: Run E2E tests
        working-directory: tools/prompt_injection_tester
        run: |
          python tests/e2e_test.py --target http://localhost:11434/api/chat --model llama3.2:1b --quick
        timeout-minutes: 10

      - name: Upload test reports
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: e2e-reports
          path: tools/prompt_injection_tester/pit_report_*.*
          retention-days: 7

  security-scan:
    name: Security Scan
    runs-on: ubuntu-latest

    steps:
      - name: Checkout code
        uses: actions/checkout@v4

      - name: Set up Python
        uses: actions/setup-python@v5
        with:
          python-version: '3.11'

      - name: Install dependencies
        working-directory: tools/prompt_injection_tester
        run: |
          python -m pip install --upgrade pip
          pip install -e .
          pip install safety bandit

      - name: Run safety check
        working-directory: tools/prompt_injection_tester
        run: |
          safety check --json
        continue-on-error: true

      - name: Run bandit security scan
        working-directory: tools/prompt_injection_tester
        run: |
          bandit -r pit/ -f json -o bandit-report.json
        continue-on-error: true

      - name: Upload security reports
        if: always()
        uses: actions/upload-artifact@v4
        with:
          name: security-reports
          path: tools/prompt_injection_tester/*-report.json
          retention-days: 30