ci cleanup

Author: ankit76

.github/workflows/ci.yml

@@ -25,10 +25,5 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install
         run: ./.github/workflows/run_ci.sh
-      - name: Install CPU JAX
-        if: matrix.python-version == '3.12'
-        run: pip install -U jax
       - name: Test with pytest
-        env:
-          PYTEST_ADDOPTS: ${{ matrix.python-version != '3.12' && '--ignore=afqmc/test' || '' }}
         run: ./.github/workflows/test.sh

.github/workflows/run_ci.sh

@@ -12,6 +12,7 @@ python -m pip install --upgrade pip
 pip install "scipy<1.16"
 pip install pytest
 pip install .
+pip install -U jax
 
 pip install trexio
 pip install mcfun

pyscf/afqmc/test/conftest.py

@@ -0,0 +1,27 @@
+import importlib.util
+import sys
+from pathlib import Path
+
+
+TEST_DIR = Path(__file__).resolve().parent
+HAS_JAX = (
+    importlib.util.find_spec("jax") is not None
+    and importlib.util.find_spec("jaxlib") is not None
+)
+
+
+def pytest_ignore_collect(collection_path, config):
+    path = Path(str(collection_path))
+    if path.parent != TEST_DIR:
+        return False
+    if path.name == "conftest.py":
+        return False
+    if path.suffix != ".py" or not path.name.startswith("test_"):
+        return False
+    if sys.version_info[:2] < (3, 10):
+        return True
+    if path.name == "test_import.py":
+        return False
+    if not HAS_JAX:
+        return True
+    return False