create run_portable_pipeline test

hjtran · Sep 17, 2023 · db01dbe · db01dbe
1 parent 38720a3
commit db01dbe
Showing 1 changed file with 12 additions and 0 deletions.
diff --git a/sdks/python/apache_beam/runners/render_test.py b/sdks/python/apache_beam/runners/render_test.py
@@ -39,6 +39,18 @@ def test_basic_graph(self):
     self.assertIn('CustomName', dot)
     self.assertEqual(dot.count('->'), 2)
 
+  def test_run_portable_pipeline(self):
+    p = beam.Pipeline()
+    _ = (
+        p | beam.Impulse() | beam.Map(lambda _: 2)
+        | 'CustomName' >> beam.Map(lambda x: x * x))
+    pipeline_proto = p.to_runner_api()
+    dot = render.RenderRunner().run_portable_pipeline(
+        pipeline_proto, default_options)
+    self.assertIn('digraph', dot)
+    self.assertIn('CustomName', dot)
+    self.assertEqual(dot.count('->'), 2)
+
   def test_side_input(self):
     p = beam.Pipeline()
     pcoll = p | beam.Impulse() | beam.FlatMap(lambda x: [1, 2, 3])