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, '')