Add notebook testing on CI

.github/workflows/notebook_tests.yml

@@ -0,0 +1,39 @@
+name: notebook tests
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  build-and-test:
+    name: Python ${{ matrix.python-version }} on ${{ matrix.os }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        python-version: ['3.8']
+        os: [ubuntu-latest]
+
+    steps:
+    - uses: actions/checkout@v2
+    - uses: actions/setup-python@v2
+      with:
+        python-version: ${{ matrix.python-version }}
+
+    - name: Build
+      run: |
+        set -xe
+        python -VV
+        pip install --upgrade pip
+        pip install -e '.[test]'
+
+    - name: Print versions
+      run: |
+        python -VV
+        python -c "import jax; print('jax', jax.__version__)"
+        python -c "import jaxlib; print('jaxlib', jaxlib.__version__)"
+
+    - name: Run notebook tests
+      run: |
+        python -m pytest --nbval docs/notebooks/ --current-env

.github/workflows/tests.yml

@@ -26,7 +26,7 @@ jobs:
       run: |
         set -xe
         python -VV
-        pip install --upgrade pip setuptools
+        pip install --upgrade pip
         pip install pytest-memray
         pip install -e '.[test]'
 
@@ -39,12 +39,12 @@ jobs:
     - name: Run fast tests
       if: ${{ matrix.test_mark == 'fast' }}
       run: |
-        pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray -m fast -n auto
+        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray -m fast -n auto
 
     - name: Run all tests
       if: ${{ matrix.test_mark == 'all' }}
       run: |
-        pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray
+        python -m pytest --cov=ott --cov-append --cov-report=xml --cov-report=term-missing --cov-config=setup.cfg --memray
 
     - name: Upload coverage
       uses: codecov/codecov-action@v3

setup.cfg

@@ -53,6 +53,7 @@ test =
     pytest
     pytest-xdist
     pytest-cov
+    nbval
     chex
 docs =
     sphinx>=4.0