internal

PiperOrigin-RevId: 448317533

Author: zwestrick

tensorflow_data_validation/api/stats_api_test.py

@@ -27,7 +27,6 @@
 import pyarrow as pa
 
 from tensorflow_data_validation.api import stats_api
-from tensorflow_data_validation.utils import statistics_io_test_util
 from tensorflow_data_validation.statistics import stats_options
 from tensorflow_data_validation.utils import io_util
 from tensorflow_data_validation.utils import stats_util
@@ -699,7 +698,7 @@ def test_write_stats_to_tfrecord_and_binary(self):
 
     output_path_binary = os.path.join(self._get_temp_dir(), 'stats.pb')
     output_path_prefix = os.path.join(self._get_temp_dir(), 'stats_shards')
-    with beam.Pipeline(runner=statistics_io_test_util.test_runner_impl()) as p:
+    with beam.Pipeline() as p:
       _ = (
           p | beam.Create([stats1, stats2])
           | stats_api.WriteStatisticsToRecordsAndBinaryFile(

tensorflow_data_validation/utils/statistics_io_impl_test.py

@@ -17,7 +17,6 @@
 from absl.testing import absltest
 import apache_beam as beam
 from tensorflow_data_validation.utils import statistics_io_impl
-from tensorflow_data_validation.utils import statistics_io_test_util
 from tensorflow_metadata.proto.v0 import statistics_pb2
 
 
@@ -32,7 +31,7 @@ def test_write_and_read_records(self):
     ]
     output_prefix = tempfile.mkdtemp() + '/statistics'
 
-    with beam.Pipeline(runner=statistics_io_test_util.test_runner_impl()) as p:
+    with beam.Pipeline() as p:
       provider = statistics_io_impl.get_io_provider('tfrecords')
       _ = (p | beam.Create(datasets) | provider.record_sink_impl(output_prefix))