def setUp(self):
"""Setup for CloudApiTest."""
super(CloudApiTest, self).setUp()
self.addCleanup(mock.patch.stopall)
self.project_id = 'project_id'
self.service_account_name = 'service_account_name'
self.api_name = 'storage'
self.session = mock.Mock()
self.mock_client = mock.Mock()
self.parent = f'projects/{self.project_id}'
self.service_account_key_file = '/path/to/service/account/key'
# Mocks for building client.
mock_build_service_client = mock.patch.object(
cloud_auth, 'build_service_client', autospec=True).start()
mock_build_service_client.return_value = self.mock_client
# Mocks for operations.
self.operation_client = mock.Mock()
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.operation = {}
self.mock_execute_request.return_value = self.operation
self.mock_client.operations.return_value = self.operation_client
self.cloud_api_utils = cloud_api.CloudApiUtils(
self.project_id, self.service_account_key_file)
def test_client_initializes_with_impersonated_service_account(
self, mock_impersonated_client):
service_account_name = '*****@*****.**'
version = 'v1beta1'
cloud_api.CloudApiUtils(self.project_id,
service_account_name=service_account_name)
mock_impersonated_client.assert_called_once_with(
'serviceusage', 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)