def test_deregister(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.cluster = CLUSTER
dereg_resp_a = copy.deepcopy(dereg_con_inst_result)
dereg_resp_a["containerInstance"][
"containerInstanceArn"
] = "arn:aws:ecs:ap-southeast-2:111111111111:container-instance/some-cluster/a99b9853b6114c87af46c7501a3a6ba8"
dereg_resp_b = copy.deepcopy(dereg_con_inst_result)
dereg_resp_b["containerInstance"][
"containerInstanceArn"
] = "arn:aws:ecs:ap-southeast-2:111111111111:container-instance/some-cluster/e338a4bea54e4e06b664530a30ae02dd"
with Stubber(boto3_clients["ecs"]) as stubber:
stubber.add_response(
method="list_container_instances",
service_response=list_cont_result,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id != {AMI}",
},
)
stubber.add_response(
method="deregister_container_instance", service_response=dereg_resp_a
)
stubber.add_response(
method="deregister_container_instance", service_response=dereg_resp_b
)
rot.deregister(session)
stubber.assert_no_pending_responses()
def test_detach_outdated_instances(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.asgs = ASGS
session.cluster = CLUSTER
desired_result = {"i-111a1a1aaa1a1aa11", "i-11111b1b1bb1bbb1b"}
with Stubber(boto3_clients["autoscaling"]) as stubber:
stubber.add_response(
method="describe_auto_scaling_groups",
service_response=des_asg_grp_result,
expected_params={"AutoScalingGroupNames": list(ASGS), "MaxRecords": 100},
)
stubber.add_response(
method="detach_instances",
service_response=det_inst_result,
expected_params={
"InstanceIds": EC2_INSTANCE_IDS,
"AutoScalingGroupName": list(ASGS)[0],
"ShouldDecrementDesiredCapacity": False,
},
)
assert rot.detach_outdated_instances(session) == desired_result
stubber.assert_no_pending_responses()
def test_drain_instances(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.cluster = CLUSTER
with Stubber(boto3_clients["ecs"]) as stubber:
stubber.add_response(
method="list_container_instances",
service_response=list_cont_result,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id != {AMI}",
},
)
stubber.add_response(
method="update_container_instances_state",
service_response=update_cont_result,
expected_params={
"cluster": CLUSTER,
"containerInstances": list_cont_result["containerInstanceArns"],
"status": "DRAINING",
},
)
rot.drain_instances(session)
stubber.assert_no_pending_responses()
def test_terminate(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.instances = ["i-11a111aa1111aa111", "i-222222222b2222222"]
with Stubber(boto3_clients["ec2"]) as stubber:
stubber.add_response(
method="terminate_instances", service_response=terminate_result
)
rot.terminate(session)
stubber.assert_no_pending_responses()
def test_find_outdated_asg(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.cluster = CLUSTER
stubs = {k: Stubber(v) for k, v in session.clients.items()}
stubs["ecs"].add_response(
method="list_container_instances",
service_response=list_cont_result,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id != {AMI}",
},
)
stubs["ecs"].add_response(
method="describe_container_instances",
service_response=desc_cont_result,
expected_params={
"cluster": CLUSTER,
"containerInstances": list_cont_result["containerInstanceArns"],
},
)
stubs["autoscaling"].add_response(
method="describe_auto_scaling_instances",
service_response=desc_asg_inst_result,
expected_params={"InstanceIds": CONTAINER_INSTANCE_IDS},
)
for stub in stubs.values():
stub.activate()
assert rot.find_outdated_asg(session) == {
"some-cluster-ECSAutoScalingGroup-ASDF1234"
}
for stub in stubs.values():
stub.assert_no_pending_responses()
stub.deactivate()
def test_can_deregister(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.asgs = ASGS
session.cluster = CLUSTER
containers = copy.deepcopy(list_cont_result)
containers["containerInstanceArns"] = [
"arn:aws:ecs:ap-southeast-2:111111111111:container-instance/some-cluster/a99b9853b6114c87af46c7501a3a6ba8"
]
success_result = copy.deepcopy(list_task_result)
success_result["taskArns"] = []
with Stubber(boto3_clients["ecs"]) as stubber:
stubber.add_response(
method="list_container_instances",
service_response=containers,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id != {AMI}",
},
)
stubber.add_response(method="list_tasks", service_response=list_task_result)
stubber.add_response(
method="list_container_instances",
service_response=containers,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id != {AMI}",
},
)
stubber.add_response(method="list_tasks", service_response=success_result)
assert rot.can_deregister(session) is False
assert rot.can_deregister(session) is True
stubber.assert_no_pending_responses()
def test_session_init():
session = rot.Session()
assert session.region == "ap-southeast-2"
assert len(session.clients) == 3
def test_can_update_service(boto3_clients):
session = rot.Session(clients=boto3_clients)
session.ami = AMI
session.asgs = ASGS
session.cluster = CLUSTER
bad_result_a = copy.deepcopy(des_asg_grp_result)
bad_result_a["AutoScalingGroups"][0]["Instances"][0]["LifecycleState"] = "Pending"
del bad_result_a["AutoScalingGroups"][0]["Instances"][1]
bad_result_b = copy.deepcopy(des_asg_grp_result)
bad_result_b["AutoScalingGroups"][0]["Instances"][0]["LifecycleState"] = "Pending"
bad_result_c = copy.deepcopy(list_cont_result)
del bad_result_c["containerInstanceArns"][1]
stubs = {k: Stubber(v) for k, v in boto3_clients.items()}
stubs["autoscaling"].add_response(
method="describe_auto_scaling_groups",
service_response=bad_result_a,
expected_params={"AutoScalingGroupNames": list(ASGS)},
)
stubs["autoscaling"].add_response(
method="describe_auto_scaling_groups",
service_response=bad_result_b,
expected_params={"AutoScalingGroupNames": list(ASGS)},
)
stubs["autoscaling"].add_response(
method="describe_auto_scaling_groups",
service_response=des_asg_grp_result,
expected_params={"AutoScalingGroupNames": list(ASGS)},
)
stubs["ecs"].add_response(
method="list_container_instances",
service_response=bad_result_c,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id == {AMI}",
},
)
stubs["autoscaling"].add_response(
method="describe_auto_scaling_groups",
service_response=des_asg_grp_result,
expected_params={"AutoScalingGroupNames": list(ASGS)},
)
stubs["ecs"].add_response(
method="list_container_instances",
service_response=list_cont_result,
expected_params={
"cluster": CLUSTER,
"filter": f"attribute:ecs.ami-id == {AMI}",
},
)
for stub in stubs.values():
stub.activate()
assert rot.can_drain_instances(session) is False
assert rot.can_drain_instances(session) is False
assert rot.can_drain_instances(session) is False
assert rot.can_drain_instances(session) is True
for stub in stubs.values():
stub.assert_no_pending_responses()
stub.deactivate()