Testing: avoid global fixture for doctests

spundun · Apr 8, 2021 · cee3af5 · cee3af5
1 parent 438b56c
commit cee3af5
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 8 deletions.
diff --git a/.github/workflows/ci-build.yaml b/.github/workflows/ci-build.yaml
@@ -141,6 +141,8 @@ jobs:
       run: |
         pip install -r docs/requirements.txt
     - name: Test documentation
+      env:
+        XLA_FLAGS: "--xla_force_host_platform_device_count=8"
       run: |
         pytest -n 1 docs
         pytest -n 1 --doctest-modules --ignore=jax/experimental/jax2tf jax
diff --git a/conftest.py b/conftest.py
@@ -15,7 +15,6 @@
 
 import jax
 import numpy
-import os
 import pytest
 
 
@@ -25,10 +24,3 @@ def add_imports(doctest_namespace):
   doctest_namespace["lax"] = jax.lax
   doctest_namespace["jnp"] = jax.numpy
   doctest_namespace["np"] = numpy
-
-
-@pytest.fixture(autouse=True)
-def spoof_devices(doctest_namespace):
-  # Set up runtime to mimic an 8-core machine
-  flags = os.environ.get('XLA_FLAGS', '')
-  os.environ['XLA_FLAGS'] = flags + " --xla_force_host_platform_device_count=8"