Update TestCaseInContext.

PiperOrigin-RevId: 616762965

Author: fineguy

tensorflow_datasets/testing/test_case.py

@@ -59,13 +59,6 @@ def setUp(self):
     # get_temp_dir is actually the same for all tests, so create a temp sub-dir.
     self.tmp_dir = tempfile.mkdtemp(dir=tf.compat.v1.test.get_temp_dir())
 
-  def assertRaisesWithPredicateMatch(self, err_type, predicate):
-    if isinstance(predicate, str):
-      predicate_fct = lambda err: predicate in str(err)
-    else:
-      predicate_fct = predicate
-    return super().assertRaisesWithPredicateMatch(err_type, predicate_fct)
-
   @contextlib.contextmanager
   def assertLogs(self, text, level="info"):
     with mock.patch.object(logging, level) as mock_log:

tensorflow_datasets/testing/test_case_in_context.py

@@ -18,10 +18,11 @@
 from collections.abc import Sequence
 import contextlib
 from typing import Any, ContextManager
-import unittest
 
+from absl.testing import absltest
 
-class TestCaseInContext(unittest.TestCase):
+
+class TestCaseInContext(absltest.TestCase):
   """Base TestCase for running tests inside the given contexts.
 
   It ensures that all contexts are entered before any test starts and are exited
@@ -47,3 +48,10 @@ def tearDownClass(cls):
     # Exit all contexts in the reverse order
     cls._stack.close()
     super().tearDownClass()
+
+  def assertRaisesWithPredicateMatch(self, err_type, predicate):  # pylint: disable=invalid-name
+    if isinstance(predicate, str):
+      predicate_fn = lambda err: predicate in str(err)
+    else:
+      predicate_fn = predicate
+    return super().assertRaisesWithPredicateMatch(err_type, predicate_fn)