github.com/apache/beam/sdks/v2@v2.48.2/python/apache_beam/runners/interactive/testing/mock_env.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  """Module of mocks to isolated the test environment for each Interactive Beam
    19  test.
    20  """
    21  
    22  import unittest
    23  import uuid
    24  from typing import Type
    25  from unittest.mock import patch
    26  
    27  from apache_beam.runners.interactive import interactive_environment as ie
    28  from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager
    29  from apache_beam.runners.interactive.interactive_environment import InteractiveEnvironment
    30  from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython
    31  
    32  
    33  def isolated_env(cls: Type[unittest.TestCase]):
    34    """A class decorator for unittest.TestCase to set up an isolated test
    35    environment for Interactive Beam."""
    36    class IsolatedInteractiveEnvironmentTest(cls):
    37      def setUp(self):
    38        self.env_patchers = []
    39        interactive_path = 'apache_beam.runners.interactive'
    40  
    41        if ie.current_env().is_interactive_ready:
    42          # Group lines of code into a notebook/IPython cell:
    43          # with self.cell:
    44          #   arbitrary python code
    45          #   ...
    46          #   arbitrary python code
    47          self.ipython_patcher = patch(
    48              'IPython.get_ipython', new_callable=mock_get_ipython)
    49          self.cell = self.ipython_patcher.start()
    50          self.env_patchers.append(self.ipython_patcher)
    51  
    52        # self.current_env IS interactive_environment.current_env() in tests.
    53        self.ie_patcher = patch(
    54            f'{interactive_path}.interactive_environment.current_env')
    55        self.m_current_env = self.ie_patcher.start()
    56        self.current_env = InteractiveEnvironment()
    57        self.m_current_env.return_value = self.current_env
    58        self.env_patchers.append(self.ie_patcher)
    59  
    60        # Patches dataproc cluster creation and deletion.
    61        self.create_cluster_patcher = patch.object(
    62            DataprocClusterManager, 'create_flink_cluster', mock_create_cluster)
    63        self.alt_create_cluster = self.create_cluster_patcher.start()
    64        self.env_patchers.append(self.create_cluster_patcher)
    65        self.delete_cluster_patcher = patch.object(
    66            DataprocClusterManager, 'cleanup')
    67        self.m_delete_cluster = self.delete_cluster_patcher.start()
    68        self.env_patchers.append(self.delete_cluster_patcher)
    69  
    70        super().setUp()
    71  
    72      def tearDown(self):
    73        super().tearDown()
    74        # Explicitly calls the cleanup instead of letting it be called at exit
    75        # when the mocks and isolated env instance are out of scope.
    76        self.current_env.cleanup()
    77        for patcher in reversed(self.env_patchers):
    78          patcher.stop()
    79  
    80    return IsolatedInteractiveEnvironmentTest
    81  
    82  
    83  def mock_create_cluster(self):
    84    """Mocks a cluster creation and populates derived fields."""
    85    self.cluster_metadata.master_url = 'test-url-' + uuid.uuid4().hex
    86    self.cluster_metadata.dashboard = 'test-dashboard'
    87  
    88  
    89  # This file contains no tests. Below lines are purely for passing lint.
    90  if __name__ == '__main__':
    91    unittest.main()