def test_fills_out_amis(iam_client_stub, ec2_client_stub): # Setup stubs to mock out boto3 stubs.configure_iam_role_default(iam_client_stub) stubs.configure_key_pair_default(ec2_client_stub) stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.configure_subnet_default(ec2_client_stub) config = helpers.load_aws_example_config_file("example-full.yaml") del config["available_node_types"]["ray.head.default"]["node_config"][ "ImageId"] del config["available_node_types"]["ray.worker.default"]["node_config"][ "ImageId"] # Pass in SG for stub to work config["head_node"]["SecurityGroupIds"] = ["sg-1234abcd"] config["worker_nodes"]["SecurityGroupIds"] = ["sg-1234abcd"] defaults_filled = bootstrap_aws(config) ami = DEFAULT_AMI.get(config.get("provider", {}).get("region")) assert defaults_filled["head_node"].get("ImageId") == ami assert defaults_filled["worker_nodes"].get("ImageId") == ami iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()
def test_iam_already_configured(iam_client_stub, ec2_client_stub): """ Checks that things work as expected when IAM role is supplied by user. """ stubs.configure_key_pair_default(ec2_client_stub) stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.configure_subnet_default(ec2_client_stub) config = helpers.load_aws_example_config_file("example-full.yaml") head_node_config = config["available_node_types"]["ray.head.default"][ "node_config"] worker_node_config = config["available_node_types"]["ray.worker.default"][ "node_config"] head_node_config["IamInstanceProfile"] = "mock_profile" # Pass in SG for stub to work head_node_config["SecurityGroupIds"] = ["sg-1234abcd"] worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"] defaults_filled = bootstrap_aws(config) filled_head = defaults_filled["available_node_types"]["ray.head.default"][ "node_config"] assert filled_head["IamInstanceProfile"] == "mock_profile" assert "IamInstanceProfile" not in defaults_filled["head_node"] iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()
def test_fills_out_amis_and_iam(iam_client_stub, ec2_client_stub, region): # Set up expected key pair for specific region region_key_pair = DEFAULT_KEY_PAIR.copy() region_key_pair["KeyName"] = DEFAULT_KEY_PAIR["KeyName"].replace( "us-west-2", region) # Setup stubs to mock out boto3 stubs.configure_iam_role_default(iam_client_stub) stubs.configure_key_pair_default(ec2_client_stub, region=region, expected_key_pair=region_key_pair) stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.configure_subnet_default(ec2_client_stub) config = helpers.load_aws_example_config_file("example-full.yaml") head_node_config = config["available_node_types"]["ray.head.default"][ "node_config"] worker_node_config = config["available_node_types"]["ray.worker.default"][ "node_config"] del head_node_config["ImageId"] del worker_node_config["ImageId"] # Pass in SG for stub to work head_node_config["SecurityGroupIds"] = ["sg-1234abcd"] worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"] config["provider"]["region"] = region defaults_filled = bootstrap_aws(config) ami = DEFAULT_AMI.get(defaults_filled.get("provider", {}).get("region")) for node_type in defaults_filled["available_node_types"].values(): node_config = node_type["node_config"] assert node_config.get("ImageId") == ami # Correctly configured IAM role assert defaults_filled["head_node"]["IamInstanceProfile"] == { "Arn": DEFAULT_INSTANCE_PROFILE["Arn"] } # Workers of the head's type do not get the IAM role. head_type = config["head_node_type"] assert ("IamInstanceProfile" not in defaults_filled["available_node_types"][head_type]) iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()
def test_subnet_given_head_and_worker_sg(iam_client_stub, ec2_client_stub): stubs.configure_iam_role_default(iam_client_stub) stubs.configure_key_pair_default(ec2_client_stub) # list a security group and a thousand subnets in different vpcs stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.describe_a_thousand_subnets_in_different_vpcs(ec2_client_stub) config = helpers.bootstrap_aws_example_config_file( "example-head-and-worker-security-group.yaml") # check that just the single subnet in the right vpc is filled assert config["head_node"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]] assert config["worker_nodes"]["SubnetIds"] == [DEFAULT_SUBNET["SubnetId"]] # expect no pending responses left in IAM or EC2 client stub queues iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()
def test_missing_keyname(iam_client_stub, ec2_client_stub): config = helpers.load_aws_example_config_file("example-full.yaml") config["auth"]["ssh_private_key"] = "/path/to/private/key" head_node_config = config["available_node_types"]["ray.head.default"][ "node_config"] worker_node_config = config["available_node_types"]["ray.worker.default"][ "node_config"] # Setup stubs to mock out boto3. Should fail on assertion after # checking KeyName/UserData. stubs.configure_iam_role_default(iam_client_stub) missing_user_data_config = copy.deepcopy(config) with pytest.raises(AssertionError): # Config specified ssh_private_key, but missing KeyName/UserData in # node configs bootstrap_aws(missing_user_data_config) # Pass in SG for stub to work head_node_config["SecurityGroupIds"] = ["sg-1234abcd"] worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"] # Set UserData for both node configs head_node_config["UserData"] = {"someKey": "someValue"} worker_node_config["UserData"] = {"someKey": "someValue"} # Stubs to mock out boto3. Should no longer fail on assertion # and go on to describe security groups + configure subnet stubs.configure_iam_role_default(iam_client_stub) stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.configure_subnet_default(ec2_client_stub) # Should work without error now that UserData is set bootstrap_aws(config) iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()
def test_log_to_cli(iam_client_stub, ec2_client_stub): config = helpers.load_aws_example_config_file("example-full.yaml") head_node_config = config["available_node_types"]["ray.head.default"][ "node_config"] worker_node_config = config["available_node_types"]["ray.worker.default"][ "node_config"] # Pass in SG for stub to work head_node_config["SecurityGroupIds"] = ["sg-1234abcd"] worker_node_config["SecurityGroupIds"] = ["sg-1234abcd"] stubs.configure_iam_role_default(iam_client_stub) stubs.configure_key_pair_default(ec2_client_stub) stubs.describe_a_security_group(ec2_client_stub, DEFAULT_SG) stubs.configure_subnet_default(ec2_client_stub) config = helpers.bootstrap_aws_config(config) # Only side-effect is to print logs to cli, called just to # check that it runs without error log_to_cli(config) iam_client_stub.assert_no_pending_responses() ec2_client_stub.assert_no_pending_responses()