def setUp(self):
"""Creates mock objects for googleapi client."""
super(CloudComposerUtilsTest, self).setUp()
self.addCleanup(mock.patch.stopall)
self.project_id = 'project_id'
self.location = cloud_composer._LOCATION
self.environment_name = 'environment_name'
self.zone = 'a'
self.mock_build_service_client = mock.patch.object(
cloud_auth, 'build_service_client', autospec=True).start()
self.mock_wait_for_operation = mock.patch.object(
utils, 'wait_for_operation', autospec=True).start()
self.mock_execute_request = mock.patch.object(utils,
'execute_request',
autospec=True).start()
self.mock_client = mock.Mock()
self.mock_build_service_client.return_value = self.mock_client
self.service_account_key_file = '/tmp/service_account_key.json'
self.composer = cloud_composer.CloudComposerUtils(
self.project_id, self.location, self.service_account_key_file)
self.operation_client = mock.Mock()
self.operation = {}
(self.mock_client.projects.return_value.locations.return_value.
operations.return_value) = self.operation_client
self.mock_execute_request.return_value = self.operation
self.http_error = errors.HttpError(mock.MagicMock(status=400), b'')
self.fully_qualified_name = (f'projects/{self.project_id}/locations/'
f'{self.location}/environments/'
f'{self.environment_name}')
self.mock_environment_client = (self.mock_client.projects.return_value.
locations.return_value.environments)
self.mock_request = mock.Mock(http.HttpRequest)
def test_client_initializes_with_impersonated_service_account(
self, mock_impersonated_client):
service_account_name = '*****@*****.**'
version = 'v1beta1'
cloud_composer.CloudComposerUtils(
project_id=self.project_id,
service_account_name=service_account_name,
version=version)
mock_impersonated_client.assert_called_once_with(
'composer', service_account_name, version)
def main() -> None:
logging.getLogger('').setLevel(logging.INFO)
args = parse_arguments()
# Create service account.
cloud_auth.create_service_account(
project_id=args.project_id,
service_account_name=_SERVICE_ACCOUNT_NAME,
role_name=_SERVICE_ACCOUNT_ROLE,
file_name=args.service_account_key_file)
# Initialize cloud util classes.
cloud_api_utils = cloud_api.CloudApiUtils(
project_id=args.project_id,
service_account_key_file=args.service_account_key_file)
cloud_composer_utils = cloud_composer.CloudComposerUtils(
project_id=args.project_id,
service_account_key_file=args.service_account_key_file)
cloud_storage_utils = cloud_storage.CloudStorageUtils(
project_id=args.project_id,
service_account_key_file=args.service_account_key_file)
# Enable required Cloud APIs.
cloud_api_utils.enable_apis(apis=_APIS_TO_BE_ENABLED)
# Create Cloud Composer environment.
cloud_composer_utils.create_environment(
environment_name=args.composer_env_name)
# Set Cloud Composer environment variables.
cloud_composer_utils.set_environment_variables(
environment_name=args.composer_env_name,
environment_variables=_COMPOSER_ENV_VARIABLES)
# Copy local DAGs and dependencies to Cloud Storage dag and plugins folders.
dags_folder_url = cloud_composer_utils.get_dags_folder(
environment_name=args.composer_env_name)
gcs_dags_path = os.path.dirname(dags_folder_url)
cloud_storage_utils.upload_directory_to_url(
source_directory_path=args.local_dags_folder,
destination_dir_url=gcs_dags_path)
# Install required Python packages on Cloud Composer environment.
cloud_composer_utils.install_python_packages(
environment_name=args.composer_env_name,
packages=_COMPOSER_PYPI_PACKAGES)