sdks/python/apache_beam/runners/render_test.py

#
# Licensed to the Apache Software Foundation (ASF) under one or more
# contributor license agreements.  See the NOTICE file distributed with
# this work for additional information regarding copyright ownership.
# The ASF licenses this file to You under the Apache License, Version 2.0
# (the "License"); you may not use this file except in compliance with
# the License.  You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# pytype: skip-file

import os
import argparse
import logging
import subprocess
import unittest
import tempfile
import pytest

import apache_beam as beam
from apache_beam.runners import render

default_options = render.RenderOptions._add_argparse_args(
    argparse.ArgumentParser()).parse_args([])


class RenderRunnerTest(unittest.TestCase):
  def test_basic_graph(self):
    p = beam.Pipeline()
    _ = (
        p | beam.Impulse() | beam.Map(lambda _: 2)
        | 'CustomName' >> beam.Map(lambda x: x * x))
    dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot()
    self.assertIn('digraph', dot)
    self.assertIn('CustomName', dot)
    self.assertEqual(dot.count('->'), 2)

  def test_render_config_validation(self):
    p = beam.Pipeline()
    _ = (
        p | beam.Impulse() | beam.Map(lambda _: 2)
        | 'CustomName' >> beam.Map(lambda x: x * x))
    pipeline_proto = p.to_runner_api()
    with pytest.raises(ValueError):
      render.RenderRunner().run_portable_pipeline(
          pipeline_proto, render.RenderOptions())

  def test_side_input(self):
    p = beam.Pipeline()
    pcoll = p | beam.Impulse() | beam.FlatMap(lambda x: [1, 2, 3])
    dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot()
    self.assertEqual(dot.count('->'), 1)
    self.assertNotIn('dashed', dot)

    _ = pcoll | beam.Map(
        lambda x, side: x * side, side=beam.pvalue.AsList(pcoll))
    dot = render.PipelineRenderer(p.to_runner_api(), default_options).to_dot()
    self.assertEqual(dot.count('->'), 3)
    self.assertIn('dashed', dot)

  def test_composite_collapse(self):
    p = beam.Pipeline()
    _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x)
    pipeline_proto = p.to_runner_api()
    renderer = render.PipelineRenderer(pipeline_proto, default_options)
    self.assertEqual(renderer.to_dot().count('->'), 8)
    create_transform_id, = [
        id
        for (id, transform) in pipeline_proto.components.transforms.items()
        if transform.unique_name == 'Create']
    renderer.update(toggle=[create_transform_id])
    self.assertEqual(renderer.to_dot().count('->'), 1)


class DotRequiringRenderingTest(unittest.TestCase):
  @classmethod
  def setUpClass(cls):
    try:
      subprocess.run(['dot', '-V'], capture_output=True, check=True)
    except FileNotFoundError:
      cls._dot_installed = False
    else:
      cls._dot_installed = True

  def setUp(self) -> None:
    if not self._dot_installed:  # type: ignore[attr-defined]
      self.skipTest('dot executable not installed')

  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()

    with tempfile.TemporaryDirectory() as tmpdir:
      svg_path = os.path.join(tmpdir, "my_output.svg")
      render.RenderRunner().run_portable_pipeline(
          pipeline_proto, render.RenderOptions(render_output=[svg_path]))
      assert os.path.exists(svg_path)

  def test_dot_well_formed(self):
    p = beam.Pipeline()
    _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x)
    pipeline_proto = p.to_runner_api()
    renderer = render.PipelineRenderer(pipeline_proto, default_options)
    # Doesn't actually look at the output, but ensures dot executes correctly.
    renderer.render_data()
    create_transform_id, = [
        id
        for (id, transform) in pipeline_proto.components.transforms.items()
        if transform.unique_name == 'Create']
    renderer.update(toggle=[create_transform_id])
    renderer.render_data()

  def test_leaf_composite_filter(self):
    p = beam.Pipeline()
    _ = p | beam.Create([1, 2, 3]) | beam.Map(lambda x: x * x)
    dot = render.PipelineRenderer(
        p.to_runner_api(),
        render.RenderOptions(['--render_leaf_composite_nodes=Create'
                              ])).to_dot()
    self.assertEqual(dot.count('->'), 1)


if __name__ == '__main__':
  logging.getLogger().setLevel(logging.INFO)
  unittest.main()