Yasuke/moto
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added DeleteComputeEnvironment and UpdateComputeEnvironment

Browse Source
This commit is contained in:
Terry Cain 2017-10-03 22:35:30 +01:00
parent a027f86cc8
commit 88a11b21ae
No known key found for this signature in database
GPG Key ID: 14D90844E4E9B9F3
5 changed files with 216 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
moto/batch/exceptions.py
View File

@ -30,3 +30,8 @@ class ValidationError(AWSError):
class InternalFailure(AWSError):
CODE = 'InternalFailure'
STATUS = 500
class ClientException(AWSError):
CODE = 'ClientException'
STATUS = 400

70
moto/batch/models.py
View File

@ -9,7 +9,7 @@ from moto.iam import iam_backends
from moto.ec2 import ec2_backends
from moto.ecs import ecs_backends
from .exceptions import InvalidParameterValueException, InternalFailure
from .exceptions import InvalidParameterValueException, InternalFailure, ClientException
from .utils import make_arn_for_compute_env
from moto.ec2.exceptions import InvalidSubnetIdError
from moto.ec2.models import INSTANCE_TYPES as EC2_INSTANCE_TYPES
@ -31,12 +31,14 @@ class ComputeEnvironment(BaseModel):
self.instances = []
self.ecs_arn = None
self.ecs_name = None
def add_instance(self, instance):
self.instances.append(instance)
def set_ecs_arn(self, arn):
def set_ecs(self, arn, name):
self.ecs_arn = arn
self.ecs_name = name
class BatchBackend(BaseBackend):
@ -75,7 +77,7 @@ class BatchBackend(BaseBackend):
self.__dict__ = {}
self.__init__(region_name)
def get_compute_environment(self, arn):
def get_compute_environment_arn(self, arn):
return self._compute_environments.get(arn)
def get_compute_environment_by_name(self, name):
@ -84,6 +86,20 @@ class BatchBackend(BaseBackend):
return comp_env
return None
def get_compute_environment(self, identifier):
"""
Get compute environment by name or ARN
:param identifier: Name or ARN
:type identifier: str
:return: Compute Environment or None
:rtype: ComputeEnvironment or None
"""
env = self.get_compute_environment_arn(identifier)
if env is None:
env = self.get_compute_environment_by_name(identifier)
return env
def describe_compute_environments(self, environments=None, max_results=None, next_token=None):
envs = set()
if environments is not None:
@ -173,7 +189,7 @@ class BatchBackend(BaseBackend):
# Should be of format P2OnDemand_Batch_UUID
cluster_name = 'OnDemand_Batch_' + str(uuid.uuid4())
ecs_cluster = self.ecs_backend.create_cluster(cluster_name)
new_comp_env.set_ecs_arn(ecs_cluster.arn)
new_comp_env.set_ecs(ecs_cluster.arn, cluster_name)
return compute_environment_name, new_comp_env.arn
@ -271,6 +287,52 @@ class BatchBackend(BaseBackend):
return instances
def delete_compute_environment(self, compute_environment_name):
if compute_environment_name is None:
raise InvalidParameterValueException('Missing computeEnvironment parameter')
compute_env = self.get_compute_environment(compute_environment_name)
if compute_env is not None:
# Pop ComputeEnvironment
self._compute_environments.pop(compute_env.arn)
# Delete ECS cluster
self.ecs_backend.delete_cluster(compute_env.ecs_name)
if compute_env.type == 'MANAGED':
# Delete compute envrionment
instance_ids = [instance.id for instance in compute_env.instances]
self.ec2_backend.terminate_instances(instance_ids)
def update_compute_environment(self, compute_environment_name, state, compute_resources, service_role):
# Validate
compute_env = self.get_compute_environment(compute_environment_name)
if compute_env is None:
raise ClientException('Compute environment {0} does not exist')
# Look for IAM role
if service_role is not None:
try:
role = self.iam_backend.get_role_by_arn(service_role)
except IAMNotFoundException:
raise InvalidParameterValueException('Could not find IAM role {0}'.format(service_role))
compute_env.service_role = role
if state is not None:
if state not in ('ENABLED', 'DISABLED'):
raise InvalidParameterValueException('state {0} must be one of ENABLED | DISABLED'.format(state))
compute_env.state = state
if compute_resources is not None:
# TODO Implement resizing of instances based on changing vCpus
# compute_resources CAN contain desiredvCpus, maxvCpus, minvCpus, and can contain none of them.
pass
return compute_env.name, compute_env.arn
available_regions = boto3.session.Session().get_available_regions("batch")
batch_backends = {region: BatchBackend(region_name=region) for region in available_regions}

34
moto/batch/responses.py
View File

@ -76,3 +76,37 @@ class BatchResponse(BaseResponse):
result = {'computeEnvironments': envs}
return json.dumps(result)
# DeleteComputeEnvironment
def deletecomputeenvironment(self):
compute_environment = self._get_param('computeEnvironment')
try:
self.batch_backend.delete_compute_environment(compute_environment)
except AWSError as err:
return err.response()
return ''
def updatecomputeenvironment(self):
compute_env_name = self._get_param('computeEnvironment')
compute_resource = self._get_param('computeResources')
service_role = self._get_param('serviceRole')
state = self._get_param('state')
try:
name, arn = self.batch_backend.update_compute_environment(
compute_environment_name=compute_env_name,
compute_resources=compute_resource,
service_role=service_role,
state=state
)
except AWSError as err:
return err.response()
result = {
'computeEnvironmentArn': arn,
'computeEnvironmentName': name
}
return json.dumps(result)

2
moto/batch/urls.py
View File

@ -8,4 +8,6 @@ url_bases = [
url_paths = {
'{0}/v1/createcomputeenvironment$': BatchResponse.dispatch,
'{0}/v1/describecomputeenvironments$': BatchResponse.dispatch,
'{0}/v1/deletecomputeenvironment': BatchResponse.dispatch,
'{0}/v1/updatecomputeenvironment': BatchResponse.dispatch,
}

109
tests/test_batch/test_batch.py
View File

@ -156,3 +156,112 @@ def test_describe_compute_environment():
)
len(resp['computeEnvironments']).should.equal(0)
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_delete_unmanaged_compute_environment():
ec2_client, iam_client, ecs_client, batch_client = _get_clients()
vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client)
compute_name = 'test_compute_env'
batch_client.create_compute_environment(
computeEnvironmentName=compute_name,
type='UNMANAGED',
state='ENABLED',
serviceRole=iam_arn
)
batch_client.delete_compute_environment(
computeEnvironment=compute_name,
)
resp = batch_client.describe_compute_environments()
len(resp['computeEnvironments']).should.equal(0)
resp = ecs_client.list_clusters()
len(resp.get('clusterArns', [])).should.equal(0)
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_delete_managed_compute_environment():
ec2_client, iam_client, ecs_client, batch_client = _get_clients()
vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client)
compute_name = 'test_compute_env'
batch_client.create_compute_environment(
computeEnvironmentName=compute_name,
type='MANAGED',
state='ENABLED',
computeResources={
'type': 'EC2',
'minvCpus': 5,
'maxvCpus': 10,
'desiredvCpus': 5,
'instanceTypes': [
't2.small',
't2.medium'
],
'imageId': 'some_image_id',
'subnets': [
subnet_id,
],
'securityGroupIds': [
sg_id,
],
'ec2KeyPair': 'string',
'instanceRole': iam_arn,
'tags': {
'string': 'string'
},
'bidPercentage': 123,
'spotIamFleetRole': 'string'
},
serviceRole=iam_arn
)
batch_client.delete_compute_environment(
computeEnvironment=compute_name,
)
resp = batch_client.describe_compute_environments()
len(resp['computeEnvironments']).should.equal(0)
resp = ec2_client.describe_instances()
resp.should.contain('Reservations')
len(resp['Reservations']).should.equal(3)
for reservation in resp['Reservations']:
reservation['Instances'][0]['State']['Name'].should.equal('terminated')
resp = ecs_client.list_clusters()
len(resp.get('clusterArns', [])).should.equal(0)
@mock_ec2
@mock_ecs
@mock_iam
@mock_batch
def test_update_unmanaged_compute_environment_state():
ec2_client, iam_client, ecs_client, batch_client = _get_clients()
vpc_id, subnet_id, sg_id, iam_arn = _setup(ec2_client, iam_client)
compute_name = 'test_compute_env'
batch_client.create_compute_environment(
computeEnvironmentName=compute_name,
type='UNMANAGED',
state='ENABLED',
serviceRole=iam_arn
)
batch_client.update_compute_environment(
computeEnvironment=compute_name,
state='DISABLED'
)
resp = batch_client.describe_compute_environments()
len(resp['computeEnvironments']).should.equal(1)
resp['computeEnvironments'][0]['state'].should.equal('DISABLED')