Source code for apache_beam.runners.interactive.testing.mock_env

# 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
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# See the License for the specific language governing permissions and
# limitations under the License.

"""Module of mocks to isolated the test environment for each Interactive Beam

import unittest
import uuid
from typing import Type
from unittest.mock import patch

from apache_beam.runners.interactive import interactive_environment as ie
from apache_beam.runners.interactive.dataproc.dataproc_cluster_manager import DataprocClusterManager
from apache_beam.runners.interactive.interactive_environment import InteractiveEnvironment
from apache_beam.runners.interactive.testing.mock_ipython import mock_get_ipython

[docs]def isolated_env(cls: Type[unittest.TestCase]): """A class decorator for unittest.TestCase to set up an isolated test environment for Interactive Beam.""" class IsolatedInteractiveEnvironmentTest(cls): def setUp(self): self.env_patchers = [] interactive_path = 'apache_beam.runners.interactive' if ie.current_env().is_interactive_ready: # Group lines of code into a notebook/IPython cell: # with self.cell: # arbitrary python code # ... # arbitrary python code self.ipython_patcher = patch( 'IPython.get_ipython', new_callable=mock_get_ipython) self.cell = self.ipython_patcher.start() self.env_patchers.append(self.ipython_patcher) # self.current_env IS interactive_environment.current_env() in tests. self.ie_patcher = patch( f'{interactive_path}.interactive_environment.current_env') self.m_current_env = self.ie_patcher.start() self.current_env = InteractiveEnvironment() self.m_current_env.return_value = self.current_env self.env_patchers.append(self.ie_patcher) # Patches dataproc cluster creation and deletion. self.create_cluster_patcher = patch.object( DataprocClusterManager, 'create_flink_cluster', mock_create_cluster) self.alt_create_cluster = self.create_cluster_patcher.start() self.env_patchers.append(self.create_cluster_patcher) self.delete_cluster_patcher = patch.object( DataprocClusterManager, 'cleanup') self.m_delete_cluster = self.delete_cluster_patcher.start() self.env_patchers.append(self.delete_cluster_patcher) super().setUp() def tearDown(self): super().tearDown() # Explicitly calls the cleanup instead of letting it be called at exit # when the mocks and isolated env instance are out of scope. self.current_env.cleanup() for patcher in reversed(self.env_patchers): patcher.stop() return IsolatedInteractiveEnvironmentTest
[docs]def mock_create_cluster(self): """Mocks a cluster creation and populates derived fields.""" self.cluster_metadata.master_url = 'test-url-' + uuid.uuid4().hex self.cluster_metadata.dashboard = 'test-dashboard'
# This file contains no tests. Below lines are purely for passing lint. if __name__ == '__main__': unittest.main()