Fix the bug when calling run() under pipeline's context manager in Beam.

PiperOrigin-RevId: 729199597

Author: tfx-copybara

tensorflow_data_validation/utils/test_util.py

@@ -445,12 +445,13 @@ def _equal(actual_results: Iterable[
 
     options = beam.options.pipeline_options.PipelineOptions(
         runtime_type_check=True)
-    with beam.Pipeline(options=options) as p:
-      result = p | beam.Create(examples) | generator.ptransform
-      util.assert_that(result, _make_result_matcher(self, expected_results))
-      pipeline_result = p.run()
-      if metrics_verify_fn:
-        metrics_verify_fn(pipeline_result.metrics())
+    p = beam.Pipeline(options=options)
+    result = p | beam.Create(examples) | generator.ptransform
+    util.assert_that(result, _make_result_matcher(self, expected_results))
+    pipeline_result = p.run()
+    pipeline_result.wait_until_finish()
+    if metrics_verify_fn:
+      metrics_verify_fn(pipeline_result.metrics())
 
 
 class CombinerFeatureStatsGeneratorTest(absltest.TestCase):