Small cleanup to CachedImplTest.test_tf_function_works_with_cache

PiperOrigin-RevId: 368225087
tensorflow · Apr 13, 2021 · dde3a30 · dde3a30
1 parent c7c9573
commit dde3a30
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/tensorflow_transform/beam/cached_impl_test.py b/tensorflow_transform/beam/cached_impl_test.py
@@ -1114,11 +1114,6 @@ def identity(x):
     self.assertMetricsCounterEqual(p2.metrics, 'saved_models_created',
                                    _ZERO_PHASE_NUM_SAVED_MODELS)
 
-    self.assertMetricsCounterEqual(p2.metrics, 'num_instances', 0)
-    self.assertMetricsCounterEqual(p2.metrics, 'cache_entries_decoded', 1)
-    self.assertMetricsCounterEqual(p2.metrics, 'cache_entries_encoded', 0)
-    self.assertMetricsCounterEqual(p2.metrics, 'saved_models_created', 1)
-
     # Modifying the tf.function contents causes cache invalidation.
     run_result = self._run_pipeline(
         feature_spec,
@@ -1135,7 +1130,8 @@ def identity(x):
     self.assertMetricsCounterEqual(p3.metrics, 'num_instances', 2)
     self.assertMetricsCounterEqual(p3.metrics, 'cache_entries_decoded', 0)
     self.assertMetricsCounterEqual(p3.metrics, 'cache_entries_encoded', 1)
-    self.assertMetricsCounterEqual(p3.metrics, 'saved_models_created', 2)
+    self.assertMetricsCounterEqual(p3.metrics, 'saved_models_created',
+                                   _SINGLE_PHASE_NUM_SAVED_MODELS)
 
   @tft_unit.named_parameters(_TF_VERSION_NAMED_PARAMETERS)
   def test_changing_constant_fails_cache(self, use_tf_compat_v1):