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)