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)
def test_service_account_is_not_created_if_already_exists(self):
"""Test service account is not created if the account already exists."""
expected_account = object()
(self.mock_service_client.return_value.get.return_value.execute.return_value
) = expected_account
mock_create = self.mock_service_client.return_value.create
actual_account = cloud_auth.create_service_account(
self.project_id, self.service_account_name, self.role_name,
self.file_name)
self.assertEqual(expected_account, actual_account)
mock_create.assert_not_called()
def test_service_account_creation(self, mock_create_sa_key, mock_set_sa_role):
"""Test that service account is created for valid input."""
expected_account = object()
(self.mock_service_client.return_value.get.return_value.execute.return_value
) = None
mock_create = self.mock_service_client.return_value.create
mock_create.return_value.execute.return_value = expected_account
actual_account = cloud_auth.create_service_account(
self.project_id, self.service_account_name, self.role_name,
self.file_name)
self.assertEqual(expected_account, actual_account)
mock_create.assert_called_once()
mock_create_sa_key.assert_called_once_with(self.project_id,
self.service_account_name,
self.file_name)
mock_set_sa_role.assert_called_once_with(self.project_id,
self.service_account_name,
self.role_name)
def test_empty_service_account_key_file_name_name_error(self):
"""Test that error is raised for invalid servic account name."""
with self.assertRaises(ValueError):
cloud_auth.create_service_account(self.project_id, self.project_id,
self.role_name, '')