Enable timeout setting for Python TestPipeline (#29646)

sdks/python/apache_beam/testing/test_pipeline.py

@@ -71,7 +71,8 @@ def __init__(
       is_integration_test=False,
       blocking=True,
       additional_pipeline_args=None,
-      display_data=None):
+      display_data=None,
+      timeout=None):
     """Initialize a pipeline object for test.
 
     Args:
@@ -96,6 +97,7 @@ def __init__(
         included when construction the pipeline options object.
       display_data (Dict[str, Any]): a dictionary of static data associated
         with this pipeline that can be displayed when it runs.
+      timeout (float optional): The time in seconds to wait for the pipeline to finish
 
     Raises:
       ValueError: if either the runner or options argument is not
@@ -107,6 +109,7 @@ def __init__(
     self.options_list = (
         self._parse_test_option_args(argv) + additional_pipeline_args)
     self.blocking = blocking
+    self.timeout = timeout
     if options is None:
       options = PipelineOptions(self.options_list)
     super().__init__(runner, options, display_data=display_data)
@@ -116,7 +119,7 @@ def run(self, test_runner_api=True):
         test_runner_api=(
             False if self.not_use_test_runner_api else test_runner_api))
     if self.blocking:
-      state = result.wait_until_finish()
+      state = result.wait_until_finish(timeout=self.timeout)
       assert state in (PipelineState.DONE, PipelineState.CANCELLED), \
           "Pipeline execution failed."