Internal change

PiperOrigin-RevId: 853799125

Author: lukebaumann

pathwaysutils/test/debug/timing_test.py

@@ -19,8 +19,8 @@
 
 from pathwaysutils.debug import timing
 
-from absl.testing import absltest
-from absl.testing import parameterized
+from google3.testing.pybase import googletest
+from google3.testing.pybase import parameterized
 
 
 class TimingTest(parameterized.TestCase):
@@ -83,4 +83,4 @@ def my_function():
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/debug/watchdog_test.py

@@ -21,8 +21,8 @@
 
 from pathwaysutils.debug import watchdog
 
-from absl.testing import absltest
-from absl.testing import parameterized
+from google3.testing.pybase import googletest
+from google3.testing.pybase import parameterized
 
 
 class WatchdogTest(parameterized.TestCase):
@@ -93,4 +93,4 @@ def test_log_thread_strack_succes(self, thread_ident, expected_log_output):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/initialize_test.py

@@ -17,8 +17,8 @@
 import jax
 from pathwaysutils import _initialize
 
-from absl.testing import absltest
-from absl.testing import parameterized
+from google3.testing.pybase import googletest
+from google3.testing.pybase import parameterized
 
 
 class InitializeTest(parameterized.TestCase):
@@ -89,4 +89,4 @@ def test_persistence_enabled(self):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/lru_cache_test.py

@@ -14,10 +14,10 @@
 
 import jax.extend
 from pathwaysutils import lru_cache
-from absl.testing import absltest
+from google3.testing.pybase import googletest
 
 
-class LruCacheTest(absltest.TestCase):
+class LruCacheTest(googletest.TestCase):
 
   def test_cache_hits(self):
     x = [100]
@@ -82,4 +82,4 @@ def f(i):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/persistence_test.py

@@ -6,10 +6,10 @@
 import numpy as np
 from pathwaysutils.persistence import helper
 
-from absl.testing import absltest
+from google3.testing.pybase import googletest
 
 
-class PersistenceTest(absltest.TestCase):
+class PersistenceTest(googletest.TestCase):
   location = "/path/to/location"
   name = "name"
   dtype = np.dtype(np.int32)
@@ -106,4 +106,4 @@ def test_get_bulk_write_request(self):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/plugin_executable_test.py

@@ -5,13 +5,13 @@
 """
 import jax
 from pathwaysutils import plugin_executable
-from absl.testing import absltest
+from google3.testing.pybase import googletest
 
 PluginExecutable = plugin_executable.PluginExecutable
 XlaRuntimeError = jax.errors.JaxRuntimeError
 
 
-class PluginExecutableTest(absltest.TestCase):
+class PluginExecutableTest(googletest.TestCase):
 
   def setUp(self):
     jax.config.update("jax_platforms", "cpu")
@@ -26,4 +26,4 @@ def test_bad_program(self):
       PluginExecutable("this is not json")
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/profiling_test.py

@@ -20,8 +20,8 @@
 from pathwaysutils import profiling
 import requests
 
-from absl.testing import absltest
-from absl.testing import parameterized
+from google3.testing.pybase import googletest
+from google3.testing.pybase import parameterized
 
 
 class ProfilingTest(parameterized.TestCase):
@@ -457,4 +457,4 @@ def test_jax_profiler_trace_calls_patched_functions(self):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()

pathwaysutils/test/proxy_backend_test.py

@@ -20,18 +20,18 @@
 from pathwaysutils import jax as pw_jax
 from pathwaysutils import proxy_backend
 
-from absl.testing import absltest
+from google3.testing.pybase import googletest
 
 
-class ProxyBackendTest(absltest.TestCase):
+class ProxyBackendTest(googletest.TestCase):
 
   def setUp(self):
     super().setUp()
     jax.config.update("jax_platforms", "proxy")
     jax.config.update("jax_backend_target", "grpc://localhost:12345")
     backend.clear_backends()
 
-  @absltest.skip("b/408025233")
+  @googletest.skip("b/408025233")
   def test_no_proxy_backend_registration_raises_error(self):
     self.assertRaises(RuntimeError, backend.backends)
 
@@ -48,4 +48,4 @@ def test_proxy_backend_registration(self):
 
 
 if __name__ == "__main__":
-  absltest.main()
+  googletest.main()