name: PR Tests
on:
  pull_request:
    types: [opened, synchronize, reopened]
    paths:
      - "nunchaku/**"
      - "src/**"
      - "tests/**"
      - "examples/**"
      - "scripts/**"
      - ".github/workflows/pr-test.yaml"
      - "!nunchaku/__version__.py"
  workflow_dispatch:
    inputs:
      ref:
        description: 'Branch, tag, or SHA to test (default: main)'
        required: false
        default: 'main'
concurrency:
  group: ${{ github.repository }}-${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}
  cancel-in-progress: true
jobs:
  run-tests:
    strategy:
      matrix:
        gpu: ['5090', '4090']
    runs-on:
      - self-hosted
      - ${{ matrix.gpu }}
      - test
    # Only run for non-draft PRs, or for workflow_dispatch events
    if: |
      (github.event_name == 'pull_request' && !github.event.pull_request.draft) ||
      (github.event_name == 'workflow_dispatch')
    env:
      COMFYUI_MODELS_ROOT: "/home/nunchaku/comfyui_models"
      NUNCHAKU_TEST_CACHE_ROOT: "/home/nunchaku/nunchaku_test_root"
    steps:
      - name: Checkout
        uses: actions/checkout@v4
        with:
          submodules: true
          ref: ${{ github.event_name == 'workflow_dispatch' && github.event.inputs.ref || '' }}
      - name: Show current commit
        run: git log -1 --oneline
      - name: Set up Python
        run: |
          which python
          echo "Setting up Python with Conda"
          conda create -n test_env python=3.13 -y
      - name: Install dependencies and build
        run: |
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          which python
          echo "Installing dependencies"
          conda install -c conda-forge gxx=11 gcc=11 libsndfile -y
          pip install uv
          uv pip install torch==2.8 torchvision==0.23 --index-url https://download.pytorch.org/whl/cu128
          uv pip install -e ".[ci]"
      - name: Setup ComfyUI
        run: |
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          which python
          cd ..
          rm -rf ComfyUI-nunchaku
          git clone -b dev https://github.com/mit-han-lab/ComfyUI-nunchaku.git
          cd ComfyUI-nunchaku
          uv pip install --torch-backend=auto -e ".[ci]"
          cd ..
          rm -rf test-workspace
          mkdir -p test-workspace
          cd test-workspace
          mkdir -p "${COMFYUI_MODELS_ROOT}"
          ln -s "${COMFYUI_MODELS_ROOT}" models
          ln -s ../ComfyUI-nunchaku/tests tests
          ln -s ../ComfyUI-nunchaku/test_data test_data
          python ../ComfyUI-nunchaku/scripts/setup_custom_nodes.py
      - name: Run ComfyUI tests
        run: |
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          pwd
          cd ../test-workspace
          pytest tests/test_workflows.py -x -vv --reruns 4 --reruns-delay 0
      - name: Run nunchaku tests
        run: |
          source $(conda info --base)/etc/profile.d/conda.sh
          conda activate test_env || { echo "Failed to activate conda env"; exit 1; }
          which python
          pytest -vv -x tests/flux/test_flux_examples.py
          pytest -vv -x tests/v1/test_examples.py
          python .github/workflows/run_all_tests.py
      - name: clean up
        if: always()
        run: |
          cd ..
          rm -rf test-workspace ComfyUI-nunchaku
