github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/interactive/display/pipeline_graph_test.py (about)

     1  #
     2  # Licensed to the Apache Software Foundation (ASF) under one or more
     3  # contributor license agreements.  See the NOTICE file distributed with
     4  # this work for additional information regarding copyright ownership.
     5  # The ASF licenses this file to You under the Apache License, Version 2.0
     6  # (the "License"); you may not use this file except in compliance with
     7  # the License.  You may obtain a copy of the License at
     8  #
     9  #    http://www.apache.org/licenses/LICENSE-2.0
    10  #
    11  # Unless required by applicable law or agreed to in writing, software
    12  # distributed under the License is distributed on an "AS IS" BASIS,
    13  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14  # See the License for the specific language governing permissions and
    15  # limitations under the License.
    16  #
    17  
    18  """Tests for apache_beam.runners.interactive.display.pipeline_graph."""
    19  # pytype: skip-file
    20  
    21  import unittest
    22  from unittest.mock import patch
    23  
    24  import apache_beam as beam
    25  from apache_beam.runners.interactive import interactive_beam as ib
    26  from apache_beam.runners.interactive import interactive_environment as ie
    27  from apache_beam.runners.interactive import interactive_runner as ir
    28  from apache_beam.runners.interactive.display import pipeline_graph
    29  from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython
    30  
    31  # pylint: disable=bad-option-value,unused-variable,possibly-unused-variable
    32  # Reason:
    33  #   Disable pylint for pipelines built for testing. Not all PCollections are
    34  #   used but they need to be assigned to variables so that we can test how
    35  #   interactive beam applies the magic around user-defined variables.
    36  
    37  
    38  # The tests need graphviz to work.
    39  @unittest.skipIf(
    40      not ie.current_env().is_interactive_ready,
    41      '[interactive] dependency is not installed.')
    42  class PipelineGraphTest(unittest.TestCase):
    43    def setUp(self):
    44      ie.new_env()
    45  
    46    def test_decoration(self):
    47      p = beam.Pipeline(ir.InteractiveRunner())
    48      # We are examining if literal `"` and trailing literal `\` are decorated
    49      # correctly.
    50      pcoll = p | '"[1]": "Create\\"' >> beam.Create(range(1000))
    51      ib.watch(locals())
    52  
    53      self.assertEqual(
    54          (
    55              'digraph G {\n'
    56              'node [color=blue, fontcolor=blue, shape=box];\n'
    57              # The py string literal from `\\\\\\"` is `\\\"` in dot and will be
    58              # rendered as `\"` because they are enclosed by `"`.
    59              '"\\"[1]\\": \\"Create\\\\\\"";\n'
    60              'pcoll [shape=circle];\n'
    61              '"\\"[1]\\": \\"Create\\\\\\"" -> pcoll;\n'
    62              '}\n'),
    63          pipeline_graph.PipelineGraph(p).get_dot())
    64  
    65    def test_get_dot(self):
    66      p = beam.Pipeline(ir.InteractiveRunner())
    67      init_pcoll = p | 'Init' >> beam.Create(range(10))
    68      squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x)
    69      cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3)
    70      ib.watch(locals())
    71  
    72      self.assertEqual((
    73          'digraph G {\n'
    74          'node [color=blue, fontcolor=blue, shape=box];\n'
    75          '"Init";\n'
    76          'init_pcoll [shape=circle];\n'
    77          '"Square";\n'
    78          'squares [shape=circle];\n'
    79          '"Cube";\n'
    80          'cubes [shape=circle];\n'
    81          '"Init" -> init_pcoll;\n'
    82          'init_pcoll -> "Square";\n'
    83          'init_pcoll -> "Cube";\n'
    84          '"Square" -> squares;\n'
    85          '"Cube" -> cubes;\n'
    86          '}\n'),
    87                       pipeline_graph.PipelineGraph(p).get_dot())
    88  
    89    @patch('IPython.get_ipython', new_callable=mock_get_ipython)
    90    def test_get_dot_within_notebook(self, cell):
    91      # Assume a mocked ipython kernel and notebook frontend have been set up.
    92      ie.current_env()._is_in_ipython = True
    93      ie.current_env()._is_in_notebook = True
    94      with cell:  # Cell 1
    95        p = beam.Pipeline(ir.InteractiveRunner())
    96        # Immediately track this local pipeline so that ipython prompts when
    97        # applying transforms will be tracked and used for labels.
    98        ib.watch(locals())
    99  
   100      with cell:  # Cell 2
   101        init_pcoll = p | 'Init' >> beam.Create(range(10))
   102  
   103      with cell:  # Cell 3
   104        squares = init_pcoll | 'Square' >> beam.Map(lambda x: x * x)
   105  
   106      with cell:  # Cell 4
   107        cubes = init_pcoll | 'Cube' >> beam.Map(lambda x: x**3)
   108  
   109      # Tracks all PCollections defined so far.
   110      ib.watch(locals())
   111      self.assertEqual((
   112          'digraph G {\n'
   113          'node [color=blue, fontcolor=blue, shape=box];\n'
   114          '"[2]: Init";\n'
   115          'init_pcoll [shape=circle];\n'
   116          '"[3]: Square";\n'
   117          'squares [shape=circle];\n'
   118          '"[4]: Cube";\n'
   119          'cubes [shape=circle];\n'
   120          '"[2]: Init" -> init_pcoll;\n'
   121          'init_pcoll -> "[3]: Square";\n'
   122          'init_pcoll -> "[4]: Cube";\n'
   123          '"[3]: Square" -> squares;\n'
   124          '"[4]: Cube" -> cubes;\n'
   125          '}\n'),
   126                       pipeline_graph.PipelineGraph(p).get_dot())
   127  
   128  
   129  if __name__ == '__main__':
   130    unittest.main()