Replace eval_saved_model/testutil with utils/test_util in several tests

PiperOrigin-RevId: 648493608

tensorflow_model_analysis/evaluators/analysis_table_evaluator_test.py

@@ -17,11 +17,11 @@
 from apache_beam.testing import util
 import tensorflow as tf
 from tensorflow_model_analysis import constants
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.evaluators import analysis_table_evaluator
+from tensorflow_model_analysis.utils import test_util
 
 
-class AnalysisTableEvaulatorTest(testutil.TensorflowModelAnalysisTest):
+class AnalysisTableEvaulatorTest(test_util.TensorflowModelAnalysisTest):
 
   def testIncludeFilter(self):
     with beam.Pipeline() as pipeline:

tensorflow_model_analysis/evaluators/evaluator_test.py

@@ -14,12 +14,12 @@
 """Test for evaluator."""
 
 import tensorflow as tf
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.evaluators import evaluator
 from tensorflow_model_analysis.extractors import extractor
+from tensorflow_model_analysis.utils import test_util
 
 
-class EvaluatorTest(testutil.TensorflowModelAnalysisTest):
+class EvaluatorTest(test_util.TensorflowModelAnalysisTest):
 
   def testVerifyEvaluatorRaisesValueError(self):
     extractors = [

tensorflow_model_analysis/evaluators/metrics_validator_test.py

@@ -15,14 +15,13 @@
 
 from absl.testing import parameterized
 import tensorflow as tf
-
 from tensorflow_model_analysis.api import types
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.evaluators import metrics_validator
 from tensorflow_model_analysis.metrics import metric_types
 from tensorflow_model_analysis.proto import config_pb2
 from tensorflow_model_analysis.proto import validation_result_pb2
 from tensorflow_model_analysis.slicer import slicer_lib as slicer
+from tensorflow_model_analysis.utils import test_util
 from google.protobuf import text_format
 
 # Tests involiving slices: (<test_name>, <slice_config> , <slice_key>)
@@ -95,8 +94,9 @@
 ], ((('feature1', 'value1'),), (('feature3', 'value3'),)))
 
 
-class MetricsValidatorTest(testutil.TensorflowModelAnalysisTest,
-                           parameterized.TestCase):
+class MetricsValidatorTest(
+    test_util.TensorflowModelAnalysisTest, parameterized.TestCase
+):
 
   def testValidateMetricsInvalidThreshold(self):
     eval_config = config_pb2.EvalConfig(

tensorflow_model_analysis/extractors/features_extractor_test.py

@@ -20,17 +20,18 @@
 import tensorflow as tf
 from tensorflow_model_analysis import constants
 from tensorflow_model_analysis.api import model_eval_lib
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.extractors import features_extractor
 from tensorflow_model_analysis.proto import config_pb2
+from tensorflow_model_analysis.utils import test_util
 from tfx_bsl.tfxio import tf_example_record
 
 from google.protobuf import text_format
 from tensorflow_metadata.proto.v0 import schema_pb2
 
 
-class FeaturesExtractorTest(testutil.TensorflowModelAnalysisTest,
-                            parameterized.TestCase):
+class FeaturesExtractorTest(
+    test_util.TensorflowModelAnalysisTest, parameterized.TestCase
+):
 
   def test_features_extractor_no_features(self):
     model_spec = config_pb2.ModelSpec()

tensorflow_model_analysis/metrics/multi_class_confusion_matrix_metrics_test.py

@@ -18,14 +18,15 @@
 from apache_beam.testing import util
 import numpy as np
 import tensorflow as tf
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.metrics import metric_types
 from tensorflow_model_analysis.metrics import metric_util
 from tensorflow_model_analysis.metrics import multi_class_confusion_matrix_metrics
+from tensorflow_model_analysis.utils import test_util
 
 
-class MultiClassConfusionMatrixMetricsTest(testutil.TensorflowModelAnalysisTest,
-                                           parameterized.TestCase):
+class MultiClassConfusionMatrixMetricsTest(
+    test_util.TensorflowModelAnalysisTest, parameterized.TestCase
+):
 
   @parameterized.named_parameters(
       {

tensorflow_model_analysis/metrics/multi_class_confusion_matrix_plot_test.py

@@ -18,14 +18,15 @@
 from apache_beam.testing import util
 import numpy as np
 import tensorflow as tf
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.metrics import metric_types
 from tensorflow_model_analysis.metrics import metric_util
 from tensorflow_model_analysis.metrics import multi_class_confusion_matrix_plot
+from tensorflow_model_analysis.utils import test_util
 
 
-class MultiClassConfusionMatrixPlotTest(testutil.TensorflowModelAnalysisTest,
-                                        parameterized.TestCase):
+class MultiClassConfusionMatrixPlotTest(
+    test_util.TensorflowModelAnalysisTest, parameterized.TestCase
+):
 
   def testMultiClassConfusionMatrixPlot(self):
     computations = (

tensorflow_model_analysis/metrics/tjur_discrimination_test.py

@@ -19,13 +19,14 @@
 from apache_beam.testing import util
 import numpy as np
 import tensorflow as tf
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.metrics import metric_util
 from tensorflow_model_analysis.metrics import tjur_discrimination
+from tensorflow_model_analysis.utils import test_util
 
 
-class TjurDisriminationTest(testutil.TensorflowModelAnalysisTest,
-                            parameterized.TestCase):
+class TjurDisriminationTest(
+    test_util.TensorflowModelAnalysisTest, parameterized.TestCase
+):
 
   @parameterized.named_parameters(
       ('coefficient_of_discrimination',

tensorflow_model_analysis/slicer/slicer_test.py

@@ -21,12 +21,12 @@
 import tensorflow as tf
 from tensorflow_model_analysis import constants
 from tensorflow_model_analysis.api import types
-from tensorflow_model_analysis.eval_saved_model import testutil
 from tensorflow_model_analysis.extractors import slice_key_extractor
 from tensorflow_model_analysis.post_export_metrics import metric_keys
 from tensorflow_model_analysis.proto import config_pb2
 from tensorflow_model_analysis.proto import metrics_for_slice_pb2
 from tensorflow_model_analysis.slicer import slicer_lib as slicer
+from tensorflow_model_analysis.utils import test_util
 from tensorflow_model_analysis.utils import util as tfma_util
 
 from google.protobuf import text_format
@@ -72,7 +72,7 @@ def wrap_fpl(fpl):
   }
 
 
-class SlicerTest(testutil.TensorflowModelAnalysisTest, parameterized.TestCase):
+class SlicerTest(test_util.TensorflowModelAnalysisTest, parameterized.TestCase):
 
   def setUp(self):
     super().setUp()

tensorflow_model_analysis/writers/writer_test.py

@@ -15,11 +15,11 @@
 
 import apache_beam as beam
 import tensorflow as tf
-from tensorflow_model_analysis.eval_saved_model import testutil
+from tensorflow_model_analysis.utils import test_util
 from tensorflow_model_analysis.writers import writer
 
 
-class WriterTest(testutil.TensorflowModelAnalysisTest):
+class WriterTest(test_util.TensorflowModelAnalysisTest):
 
   def testWriteIgnoresMissingKeys(self):
     with beam.Pipeline() as pipeline: