test.yaml

name: Test

on:
  # Trigger the workflow on push or pull request, but only on main branch
  push:
    branches:
      - main
  pull_request:
    branches:
      - main

permissions:
  contents: read  # to fetch code

concurrency:
  group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
  cancel-in-progress: true

# TODO(jakevdp): add testing on macOS-11 and windows-2019 when grain supports them,
# or alternatively run a subset of tests without grain.
jobs:
  build:
    name: Pinned packages (${{ matrix.os }} Python ${{ matrix.python-version }})
    runs-on: ${{ matrix.os }}
    timeout-minutes: 10
    strategy:
      matrix:
        os: ["ubuntu-latest"]
        python-version: ["3.10", "3.11", "3.12"]

    steps:
    - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683  # v4.2.2
    - name: Set up Python ${{ matrix.python-version }}
      uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b  # v5.3.0
      with:
        python-version: ${{ matrix.python-version }}
    - name: Install dependencies
      run: |
        python -m pip install --upgrade pip
        pip install .[dev,tfds,grain]
    - name: Run tests
      run: |
        pytest -n auto jax_ai_stack