Add python test to reproduce a failed dofn.

Author: shunping

sdks/python/apache_beam/runners/portability/prism_runner_test.py

@@ -21,7 +21,6 @@
 import os.path
 import queue
 import shlex
-import threading
 import time
 import typing
 import unittest
@@ -189,6 +188,7 @@ def create_options(self):
     options.view_as(StandardOptions).streaming = self.streaming
     options.view_as(
         TypeOptions).allow_unsafe_triggers = self.allow_unsafe_triggers
+    options.view_as(PortableOptions).job_server_timeout = 10
     return options
 
   # Can't read host files from within docker, read a "local" file there.
@@ -300,6 +300,25 @@ def test_after_count_trigger_streaming(self):
                   ('B-3', {10, 15, 16}),
               ])))
 
+  def test_dofn_failure_clean_exit(self):
+    class FailDoFn(beam.DoFn):
+      def process(self, element):
+        raise ValueError("Failing as intended")
+
+    class BlockDoFn(beam.DoFn):
+      def process(self, element):
+        time.sleep(1000)
+        yield element
+
+    with self.assertRaisesRegex(Exception, "Failing as intended"):
+      with self.create_pipeline() as p:
+        imp = p | beam.Create([1, 2])
+        # Ensure the steps are not fused (otherwise siblings are run sequentially
+        # in a single thread, making execution order dependent on internal map
+        # traversal). Reshuffle acts as a fusion break so they run in parallel.
+        _ = imp | 'ReshuffleBlock' >> beam.Reshuffle() | 'Block' >> beam.ParDo(BlockDoFn())
+        _ = imp | 'ReshuffleFail' >> beam.Reshuffle() | 'Fail' >> beam.ParDo(FailDoFn())
+
 
 class PrismJobServerTest(unittest.TestCase):
   def setUp(self) -> None: