def test_get_single_waiter_config(self): single_waiter = { 'description': 'Waiter description', 'operation': 'HeadBucket', 'delay': 5, 'maxAttempts': 20, 'acceptors': [ { 'state': 'success', 'matcher': 'status', 'expected': 200 }, { 'state': 'retry', 'matcher': 'status', 'expected': 404 }, ], } waiters = { 'version': 2, 'waiters': { 'BucketExists': single_waiter, } } model = WaiterModel(waiters) config = model.get_waiter('BucketExists') self.assertEqual(config.operation, 'HeadBucket')
def test_get_waiter_does_not_exist(self): waiters = { 'version': 2, 'waiters': {} } model = WaiterModel(waiters) with self.assertRaises(ValueError): model.get_waiter('UnknownWaiter')
def test_waiter_names(self): waiters = { 'version': 2, 'waiters': { 'BarWaiter': {}, 'FooWaiter': {}, } } self.assertEqual( WaiterModel(waiters).waiter_names, ['BarWaiter', 'FooWaiter'])
def test_get_single_waiter_config(self): single_waiter = { 'description': 'Waiter description', 'operation': 'HeadBucket', 'delay': 5, 'maxAttempts': 20, 'acceptors': [ {'state': 'success', 'matcher': 'status', 'expected': 200}, {'state': 'retry', 'matcher': 'status', 'expected': 404}, ], } waiters = { 'version': 2, 'waiters': { 'BucketExists': single_waiter, } } model = WaiterModel(waiters) config = model.get_waiter('BucketExists') self.assertEqual(config.operation, 'HeadBucket')
def test_add_waiters_no_waiter_names(self): self.session.get_waiter_model.return_value = WaiterModel({ 'version': 2, # No waiters are specified. 'waiters': {} }) command_table = {} add_waiters(command_table, self.session, self.command_object) # Make sure that no wait command was added since the service object # has no waiters. self.assertEqual(command_table, {})
def setUp(self): self.model = WaiterModel({ 'version': 2, 'waiters': { 'Foo': { 'operation': 'foo', 'maxAttempts': 1, 'delay': 1, 'acceptors': [], } } }) self.service_object = mock.Mock() self.cmd = WaitCommand(self.model, self.service_object)
def get_waiter_model(self, service, api_version=None): """ Get the waiter model for the service """ service = os.path.join('aws', service) model_version = self.loader.determine_latest(service, api_version) # Some wierd formatting required to get the name of the model # correct. Right now this is returned: YYYY-MM-DD.api # We need: YYYY-MM-DD model_version = ''.join(model_version.split('.')[:-1]) waiter_model = model_version + '.waiters' return WaiterModel(self.loader.load_data(waiter_model))
def setUp(self): self.waiter_config = { 'version': 2, 'waiters': { 'WaiterName': { 'operation': 'Foo', 'delay': 1, 'maxAttempts': 1, 'acceptors': [], }, }, } self.waiter_model = WaiterModel(self.waiter_config)
def setUp(self): self.service_model = mock.Mock() self.session = mock.Mock() self.command_object = mock.Mock() self.command_object.service_model = self.service_model # Set up the mock session. self.session.get_waiter_model.return_value = WaiterModel({ 'version': 2, 'waiters': { 'FooExists': {}, } })
def tgw_attachment_waiter(self, desired_state: AttachmentState, attachment_id: str) -> None: """[summary] Args: desired_state (str): desired state of the tgw attachment attachment_id (str): attachment-id """ delay = 10 max_attempts = 15 waiter_name = "TGWAttachmentInPendingAcceptance" waiter_config = { "version": 2, "waiters": { "TGWAttachmentInPendingAcceptance": { "operation": "DescribeTransitGatewayPeeringAttachments", "delay": delay, "maxAttempts": max_attempts, "acceptors": [ { "matcher": "path", "expected": desired_state.value, "argument": "TransitGatewayPeeringAttachments[0].State", "state": "success", }, { "matcher": "path", "expected": AttachmentState.FAILED.value, "argument": "TransitGatewayPeeringAttachments[0].State", "state": "failure", }, ], } }, } waiter_model = WaiterModel(waiter_config) custom_waiter = create_waiter_with_client(waiter_name, waiter_model, self.ec2_client) try: custom_waiter.wait(TransitGatewayAttachmentIds=[attachment_id]) except WaiterError as err: self.logger.error(str(err)) raise
def _load_prefect_waiter(boto_client: "boto3.client", client_str: str, waiter_name: str): """ Load a custom waiter from the ./waiters directory. """ try: # Instantiate waiter from accompanying client json file with pkg_resources.open_text(waiters, f"{client_str}.json") as handle: waiter_model = WaiterModel(json.load(handle)) return create_waiter_with_client(waiter_name, waiter_model, boto_client) except Exception as err: raise ValueError( f"Unable to load waiter '{waiter_name}' for AWS client '{client_str}'." ) from err
def setUp(self): self.waiter_config = { 'version': 2, 'waiters': { 'WaiterName': { 'operation': 'Foo', 'delay': 1, 'maxAttempts': 1, 'acceptors': [], }, }, } self.waiter_model = WaiterModel(self.waiter_config) self.service_json_model = { 'metadata': { 'serviceFullName': 'Amazon MyService' }, 'operations': { 'Foo': { 'name': 'Foo', 'input': { 'shape': 'FooInputOutput' }, 'output': { 'shape': 'FooInputOutput' } } }, 'shapes': { 'FooInputOutput': { 'type': 'structure', 'members': { 'bar': { 'shape': 'String', 'documentation': 'Documents bar' } } }, 'String': { 'type': 'string' } } } self.service_model = ServiceModel(self.service_json_model, 'myservice') self.client = mock.Mock() self.client.meta.service_model = self.service_model
def gen_available_waiter(WAITER_ID, operation, delay, maxAttempts, path): return WaiterModel({ 'version': 2, 'waiters': { WAITER_ID: { 'operation': operation, 'delay': delay, 'maxAttempts': maxAttempts, 'acceptors': [{ 'state': 'success', 'matcher': 'path', 'argument': f"{path} == 'available'", 'expected': True }] } } })
def handle_waiters(client, client_name, class_name, service_name, service_path, sidebar_lines): waiter_config = client._get_waiter_config() waiter_model = WaiterModel( waiter_config) if 'waiters' in waiter_config else None if not waiter_model: return waiters_path = f'{service_path}/waiters' sidebar_lines.append(f' - [Waiters]({waiters_path})') docs_waiters_path = f'docs/{waiters_path}.md' waiter_names = waiter_model.waiter_names example_waiter_name = waiter_names[0] waiter_list_items = create_waiter_index(docs_waiters_path, client_name, service_name, example_waiter_name) for name in waiter_names: handle_waiter(class_name, client_name, name, service_path, waiter_list_items, waiter_model, waiters_path) write_lines(docs_waiters_path, waiter_list_items)
def delete_fargate_profile( profile_name: str, cluster_name: str, ) -> None: _logger.debug(f"Deleting EKS Fargate Profile: {profile_name}") if describe_fargate_profile(profile_name=profile_name, cluster_name=cluster_name) is None: _logger.debug(f"EKS Fargate Profile not found: {profile_name}") return eks_client = boto3_client("eks") eks_client.delete_fargate_profile( fargateProfileName=profile_name, clusterName=cluster_name, ) waiter_model = WaiterModel(WAITER_CONFIG) waiter = create_waiter_with_client("FargateProfileDeleted", waiter_model, eks_client) waiter.wait(fargateProfileName=profile_name, clusterName=cluster_name) _logger.debug(f"Deleted EKS Fargate Profile: {profile_name}")
def get_class_output(client_name): method_signatures = [] shapes_in_classes = [] client = boto3.client(client_name) class_name = type(client).__name__ service_model = client._service_model waiter_config = client._get_waiter_config() waiter_model = WaiterModel( waiter_config) if 'waiters' in waiter_config else None try: paginator_model = botocore.session.get_session().get_paginator_model( client_name) except botocore.exceptions.UnknownServiceError: paginator_model = None # meaning it probably doesn't have paginators for name in service_model.operation_names: method_signatures.append( get_method_signature(service_model, name, shapes_in_classes, class_name)) return get_class_signature(client_name, class_name, service_model.documentation, method_signatures, shapes_in_classes, waiter_model, paginator_model)
def _create_delete_traffic_session_waiter(ec2_client): delete_session_model = WaiterModel({ "version": 2, "waiters": { "TrafficMirrorDeleted": { "delay": 15, "operation": "DescribeTrafficMirrorSessions", "maxAttempts": 40, "acceptors": [{ "matcher": "error", "expected": "InvalidTrafficMirrorSessionId.NotFound", "state": "success", }], } }, }) delete_session_waiter = create_waiter_with_client( "TrafficMirrorDeleted", delete_session_model, ec2_client) return delete_session_waiter
def setUp(self): self.service_object = mock.Mock() # Create some waiters. self.model = WaiterModel({ 'version': 2, 'waiters': { 'InstanceRunning': { 'description': 'my waiter description', 'delay': 1, 'maxAttempts': 10, 'operation': 'MyOperation', }, 'BucketExists': { 'description': 'my waiter description', 'operation': 'MyOperation', 'delay': 1, 'maxAttempts': 10, } } }) self.waiter_builder = WaiterStateCommandBuilder( self.model, self.service_object)
def get_waiter_model(self, service, api_version=None): """Get the waiter model for the service.""" return WaiterModel( self.loader.load_service_model(service, type_name='waiters-2', api_version=api_version))
def get_waiter_model(self, service, api_version=None): """Get the waiter model for the service.""" with mock.patch('botocore.loaders.Loader.list_available_services', return_value=[service]): return WaiterModel(self.loader.load_service_model( service, type_name='waiters-2', api_version=api_version))
def create_environment(self, auth_credentials: Dict, user: User, environment: Environment): org_client = self._get_client("organizations") # Create an account. Requires organizations:CreateAccount permission account_request = org_client.create_account( Email=user.email, AccountName=uuid4().hex, IamUserAccessToBilling="ALLOW") # Configuration for our CreateAccount Waiter. # A waiter is a boto3 helper which can be configured to poll a given status # endpoint until it succeeds or fails. boto3 has many built in waiters, but none # for the organizations service so we're building our own here. waiter_config = { "version": 2, "waiters": { "AccountCreated": { "operation": "DescribeCreateAccountStatus", "delay": 20, "maxAttempts": self.MAX_CREATE_ACCOUNT_ATTEMPTS, "acceptors": [ { "matcher": "path", "expected": "SUCCEEDED", "argument": "CreateAccountStatus.State", "state": "success", }, { "matcher": "path", "expected": "IN_PROGRESS", "argument": "CreateAccountStatus.State", "state": "retry", }, { "matcher": "path", "expected": "FAILED", "argument": "CreateAccountStatus.State", "state": "failure", }, ], } }, } waiter_model = WaiterModel(waiter_config) account_waiter = create_waiter_with_client("AccountCreated", waiter_model, org_client) try: # Poll until the CreateAccount request either succeeds or fails. account_waiter.wait( CreateAccountRequestId=account_request["CreateAccountStatus"] ["Id"]) except WaiterError: # TODO: Possible failure reasons: # 'ACCOUNT_LIMIT_EXCEEDED'|'EMAIL_ALREADY_EXISTS'|'INVALID_ADDRESS'|'INVALID_EMAIL'|'CONCURRENT_ACCOUNT_MODIFICATION'|'INTERNAL_FAILURE' raise EnvironmentCreationException(environment.id, "Failed to create account.") # We need to re-fetch this since the Waiter throws away the success response for some reason. created_account_status = org_client.describe_create_account_status( CreateAccountRequestId=account_request["CreateAccountStatus"] ["Id"]) account_id = created_account_status["CreateAccountStatus"]["AccountId"] return account_id
def lambda_handler(event, context): # Set base parameters for each simulation job. Order from Lambda event, then environment variables. if 'simulationJobParams' in event: if 'vpcConfig' in event['simulationJobParams']: sim_job_params['vpcConfig'] = sim_job_params['vpcConfig'] if 'iamRole' in event['simulationJobParams']: sim_job_params['iamRole'] = sim_job_params['iamRole'] if 'outputLocation' in event['simulationJobParams']: sim_job_params['outputLocation'] = sim_job_params['outputLocation'] if 'simulationApplicationArn' in event: app_arn = event['simulationApplicationArn'] if 'serverIP' in event: private_ip = event['serverIP'] else: # Launch Server. server_app_params = create_application_config(event['server'], True, 'localhost') server_job_params = deepcopy(sim_job_params) server_job_params['simulationApplications'].append(server_app_params) server_job_response = robomaker.create_simulation_job( iamRole=server_job_params["iamRole"], maxJobDurationInSeconds=server_job_params["maxJobDurationInSeconds"], simulationApplications=[server_app_params], vpcConfig=server_job_params["vpcConfig"], loggingConfig=server_job_params["loggingConfig"], outputLocation=server_job_params["outputLocation"] ) # Wait for server to be available. waiter_name = 'SimJobCreated' waiter_model = WaiterModel(waiter_config) custom_waiter = create_waiter_with_client(waiter_name, waiter_model, robomaker) custom_waiter.wait(job=server_job_response['arn']) desc_result = robomaker.describe_simulation_job( job = server_job_response['arn'] ) private_ip = desc_result['networkInterface']['privateIpAddress'] # Launch multiple robot batches. batch_job_requests = [] client_app_params = {} client_job_params = {} for robot in event['robots']: client_app_params[robot['name']] = create_application_config(robot, False, private_ip) client_job_params[robot['name']] = deepcopy(sim_job_params) client_job_params[robot['name']]['simulationApplications'].append(client_app_params[robot['name']]) batch_job_requests.append(client_job_params[robot['name']]) response = robomaker.start_simulation_job_batch( batchPolicy={ 'timeoutInSeconds': DEFAULT_MAX_DURATION, 'maxConcurrency': len(event['robots']) }, createSimulationJobRequests=batch_job_requests, tags = { 'launcher': 'multi_robot_fleet' }) return { 'statusCode': 200, 'body': response['arn'] }
def get_waiter_model(config_file): waiter_dir = Path(__file__).parent.parent.joinpath("waiters") with open(waiter_dir.joinpath(config_file)) as f: config = json.load(f) return WaiterModel(config)
def delete_fargate_profile( profile_name: str, cluster_name: str, ) -> None: _logger.debug(f"Deleting EKS Fargate Profile: {profile_name}") if describe_fargate_profile(profile_name=profile_name, cluster_name=cluster_name) is None: _logger.debug(f"EKS Fargate Profile not found: {profile_name}") return eks_client = boto3_client("eks") try: timeout = 0 _logger.debug(f"Checking {cluster_name} cluster current state") res = eks_client.describe_cluster(name=cluster_name) eks_status = res.get("cluster").get("status") _logger.debug(f"{cluster_name} current state: {eks_status}") while eks_status != "ACTIVE" and timeout < 15: _logger.debug( f"EKS is an {eks_status} state. Retrying in 1min - attempt {timeout}" ) time.sleep(60) timeout += 1 res = eks_client.describe_cluster(name=cluster_name) eks_status = res.get("cluster").get("status") timeout = 0 _logger.debug(f"Checking fargate profile {profile_name} current state") fargate_profile_res = eks_client.describe_fargate_profile( clusterName=cluster_name, fargateProfileName=profile_name) fargate_profile_status = fargate_profile_res.get("fargateProfile").get( "status") _logger.debug( f"{profile_name} current state: {fargate_profile_status}") while fargate_profile_status != "ACTIVE" and timeout < 15: _logger.debug( f"Fargate profile is in {fargate_profile_status} state. Retrying in 1min - attempt {timeout}" ) time.sleep(60) timeout += 1 fargate_profile_res = eks_client.describe_fargate_profile( clusterName=cluster_name, fargateProfileName=profile_name) fargate_profile_status = fargate_profile_res.get( "fargateProfile").get("status") _logger.debug(f"Deleting fargate profile {profile_name}") eks_client.delete_fargate_profile( fargateProfileName=profile_name, clusterName=cluster_name, ) except eks_client.exceptions.ResourceNotFoundException as err: _logger.debug(err) waiter_model = WaiterModel(WAITER_CONFIG) waiter = create_waiter_with_client("FargateProfileDeleted", waiter_model, eks_client) waiter.wait(fargateProfileName=profile_name, clusterName=cluster_name) _logger.debug(f"Deleted EKS Fargate Profile: {profile_name}")
def __init__(self, client): waiter_json_filename = os.path.join(utils.__path__[0], 'cfn-waiters-2.json') with open(waiter_json_filename, 'r') as waiter_json_file: self.waiter_json_model = json.load(waiter_json_file) self.waiter_model = WaiterModel(self.waiter_json_model) self.waiter = create_waiter_with_client('StackAvailable', self.waiter_model, client.meta.client)
def test_unsupported_waiter_version(self): waiters = {'version': 1, 'waiters': {}} with self.assertRaises(WaiterConfigError): WaiterModel(waiters)
def test_waiter_version(self): self.assertEqual(WaiterModel({'version': 2, 'waiters': {}}).version, 2)
def test_get_waiter_does_not_exist(self): waiters = {'version': 2, 'waiters': {}} model = WaiterModel(waiters) with self.assertRaises(ValueError): model.get_waiter('UnknownWaiter')
def run( self, client: str = None, waiter_name: str = None, waiter_definition: dict = None, waiter_kwargs: dict = None, credentials: str = None, ): """ Task for waiting on a long-running AWS job. Uses the underlying boto3 waiter functionality. Args: - client (str): The AWS client on which to wait (e.g., 'batch', 'ec2', etc) - waiter_name (str, optional): The name of the waiter to instantiate. Can be a boto-supported waiter or one of prefect's custom waiters. Currently, prefect offers three additional waiters for AWS Batch: `"JobExists"` waits for a job to be instantiated, `"JobRunning"` waits for a job to start running, and `"JobComplete"` waits for a job to finish. You can find the definitions for all prefect-defined waiters [here](https://github.com/PrefectHQ/prefect/tree/master/src/prefect/tasks/aws/waiters). # noqa You may also use a custom waiter name, if you supply an accompanying waiter definition dict. - waiter_definition (dict, optional): A valid custom waiter model, as a dict. Note that if you supply a custom definition, it is assumed that the provided 'waiter_name' is contained within the waiter definition dict. - waiter_kwargs (dict, optional): Arguments to pass to the `waiter.wait(...)` method. Will depend upon the specific waiter being called. - credentials (dict, optional): your AWS credentials passed from an upstream Secret task; this Secret must be a JSON string with two keys: `ACCESS_KEY` and `SECRET_ACCESS_KEY` which will be passed directly to `boto3`. If not provided here or in context, `boto3` will fall back on standard AWS rules for authentication. """ if not client: raise ValueError("An AWS client string must be provided.") if not waiter_name: raise ValueError("A waiter name must be provided.") if not waiter_kwargs: waiter_kwargs = {} boto_client = get_boto_client(client, credentials=credentials, **self.boto_kwargs) if waiter_definition: # Use user-provided waiter definition waiter_model = WaiterModel(waiter_definition) waiter = create_waiter_with_client(waiter_name, waiter_model, boto_client) else: # Use either boto-provided or prefect-provided waiter if waiter_name in boto_client.waiter_names: waiter = boto_client.get_waiter(waiter_name) else: waiter = self._load_prefect_waiter(boto_client, client, waiter_name) try: waiter.wait(**waiter_kwargs) except WaiterError as e: raise FAIL( f"AWS {client} waiter '{waiter_name}' failed with: {str(e)}" ) from e