Add cleanup to restore JAX config and environment variables in tests.

PiperOrigin-RevId: 875255664

Author: lukebaumann

pathwaysutils/test/initialize_test.py

@@ -22,6 +22,19 @@
 
 class InitializeTest(parameterized.TestCase):
 
+  def setUp(self):
+    super().setUp()
+
+    orig_jax_platforms = getattr(jax.config, "jax_platforms", None)
+    self.addCleanup(jax.config.update, "jax_platforms", orig_jax_platforms)
+
+    self._orig_environ = os.environ.copy()
+    self.addCleanup(self._restore_environ)
+
+  def _restore_environ(self):
+    os.environ.clear()
+    os.environ.update(self._orig_environ)
+
   def test_first_initialize(self):
     jax.config.update("jax_platforms", "proxy")
     _initialize._initialization_count = 0

pathwaysutils/test/persistence_test.py

@@ -16,9 +16,13 @@ class PersistenceTest(absltest.TestCase):
   timeout = datetime.timedelta(seconds=3)
 
   def setUp(self):
-    jax.config.update("jax_platforms", "cpu")
     super().setUp()
 
+    orig_jax_platforms = getattr(jax.config, "jax_platforms", None)
+    self.addCleanup(jax.config.update, "jax_platforms", orig_jax_platforms)
+
+    jax.config.update("jax_platforms", "cpu")
+
   def test_get_read_request(self):
     devices = jax.devices()
     # 1d sharding

pathwaysutils/test/plugin_executable_test.py

@@ -14,9 +14,13 @@
 class PluginExecutableTest(absltest.TestCase):
 
   def setUp(self):
-    jax.config.update("jax_platforms", "cpu")
     super().setUp()
 
+    orig_jax_platforms = getattr(jax.config, "jax_platforms", None)
+    self.addCleanup(jax.config.update, "jax_platforms", orig_jax_platforms)
+
+    jax.config.update("jax_platforms", "cpu")
+
   def test_bad_json_program(self):
     with self.assertRaisesRegex(XlaRuntimeError, "INVALID_ARGUMENT"):
       PluginExecutable('{"printTextRequest":{"badParamName":"foo"}}')

pathwaysutils/test/proxy_backend_test.py

@@ -26,6 +26,15 @@ class ProxyBackendTest(absltest.TestCase):
 
   def setUp(self):
     super().setUp()
+    orig_jax_platforms = getattr(jax.config, "jax_platforms", None)
+    orig_jax_backend_target = getattr(jax.config, "jax_backend_target", None)
+
+    self.addCleanup(jax.config.update, "jax_platforms", orig_jax_platforms)
+    self.addCleanup(
+        jax.config.update, "jax_backend_target", orig_jax_backend_target
+    )
+    self.addCleanup(backend.clear_backends)
+
     jax.config.update("jax_platforms", "proxy")
     jax.config.update("jax_backend_target", "grpc://localhost:12345")
     backend.clear_backends()