test: use nox

Signed-off-by: nstarman <[email protected]>

.github/workflows/ci.yml

@@ -132,13 +132,29 @@ jobs:
     needs: [smoke]
     steps:
       - uses: actions/checkout@v4
+        with:
+          fetch-depth: 0
       - uses: actions/setup-python@v5
         with:
           python-version: "3.12"
+          allow-prereleases: true
+      # - name: Install dependencies
+      #   run: python -m pip install '.[test]'
+      # - name: Run benchmarks
+      #   uses: CodSpeedHQ/action@v3
+      #   with:
+      #     token: ${{ secrets.CODSPEED_TOKEN }}
+      #     run: python -m pytest tests/benchmark --codspeed
+
+      - name: Install Nox
+        run: pip install nox
+
       - name: Install dependencies
-        run: python -m pip install '.[test]'
-      - name: Run benchmarks
+        run: nox --sessions benchmark --install-only
+
+      - name: Run the action
         uses: CodSpeedHQ/action@v3
         with:
+          run:
+            nox --sessions benchmark --reuse-existing-virtualenvs --no-install
           token: ${{ secrets.CODSPEED_TOKEN }}
-          run: pytest tests/benchmark --codspeed

noxfile.py

@@ -39,7 +39,7 @@ def pylint(session: nox.Session) -> None:
 @nox.session
 def tests(session: nox.Session) -> None:
     """Run the unit and regular tests."""
-    session.install("-e", ".[test]")
+    session.install(".[test]")
     os.environ["GALAX_ENABLE_RUNTIME_TYPECHECKS"] = "1"  # TODO: set in a better way
     session.run("pytest", *session.posargs)
 
@@ -124,3 +124,10 @@ def build(session: nox.Session) -> None:
 
     session.install("build")
     session.run("python", "-m", "build")
+
+
+@nox.session
+def benchmark(session: nox.Session) -> None:
+    """Run the benchmark tests."""
+    session.install(".[test]")
+    session.run("python", "-m", "pytest", "tests/benchmark", "--codspeed")

tests/benchmark/__init__.py

@@ -0,0 +1,5 @@
+"""Benchmark tests."""
+
+import jax
+
+jax.config.update("jax_enable_x64", True)  # noqa: FBT003