diff --git a/ardere/aws.py b/ardere/aws.py index b11104e..31982db 100644 --- a/ardere/aws.py +++ b/ardere/aws.py @@ -1,5 +1,4 @@ """AWS Helper Classes""" -import json import logging import os import time @@ -9,13 +8,7 @@ import boto3 import botocore from concurrent.futures import ThreadPoolExecutor -from typing import ( - Any, - Dict, - List, - Optional, - Tuple -) # noqa +from typing import (Any, Dict, List, Optional, Tuple) # noqa logger = logging.getLogger() logger.setLevel(logging.INFO) @@ -314,7 +307,7 @@ def request_instances(self, instances, security_group_ids, def locate_metrics_container_ip(self): # type: () -> Tuple[Optional[str], Optional[str]] """Locates the metrics container IP and container instance arn - + Returns a tuple of (public_ip, container_arn) """ diff --git a/ardere/scripts/__init__.py b/ardere/scripts/__init__.py index 8fa2136..792d600 100644 --- a/ardere/scripts/__init__.py +++ b/ardere/scripts/__init__.py @@ -1,2 +1 @@ # - diff --git a/ardere/step_functions.py b/ardere/step_functions.py index d4e7b08..445231e 100644 --- a/ardere/step_functions.py +++ b/ardere/step_functions.py @@ -1,4 +1,3 @@ -import json import logging import os import re @@ -7,7 +6,6 @@ import boto3 import botocore -import requests import toml from marshmallow import ( Schema, @@ -279,7 +277,8 @@ def ensure_metric_sources_created(self): raise CreatingMetricSourceException("Started metric creation") if not self.ecs.has_finished_metric_creation(): - raise CreatingMetricSourceException("Metric creation still running") + raise CreatingMetricSourceException("Metric creation still " + "running") metric_ip = self.event["influxdb_private_ip"] self.event["grafana_dashboard"] = "http://{}:3000".format(metric_ip) @@ -368,15 +367,23 @@ def check_drained(self): """ client = self.boto.client('ecs') - actives = len( - client.list_container_instances( - cluster=self.event["ecs_name"], - maxResults=1, - status="ACTIVE", - ).get('containerInstanceArns', [])) - if actives: + actives = client.list_container_instances( + cluster=self.event["ecs_name"], + maxResults=1, + status="ACTIVE", + ).get('containerInstanceArns', []) + # filter out metric servers + if self.event["metrics_options"]["enabled"]: + metrics = self.ecs.locate_metrics_service() + if metrics: + metrics_arn = metrics.get("serviceArn") + try: + actives.remove(metrics_arn) + except ValueError: + pass + if len(actives): raise UndrainedInstancesException( - "Still {} active.".format(actives)) + "Still active: {}.".format(actives)) draining = len( client.list_container_instances( cluster=self.event["ecs_name"], @@ -385,5 +392,5 @@ def check_drained(self): ).get('containerInstanceArns', [])) if draining: raise UndrainedInstancesException( - "Still {} draining.".format(draining)) + "Still draining: {}.".format(draining)) return self.event diff --git a/tests/test_step_functions.py b/tests/test_step_functions.py index 99a78b5..669fc91 100644 --- a/tests/test_step_functions.py +++ b/tests/test_step_functions.py @@ -300,6 +300,7 @@ def test_drain_check_active(self): mock_client.list_container_instances.return_value = { 'containerInstanceArns': [ 'Some-Arn-01234567890', + 'Metric-Arn-01234567890', ], "nextToken": "token-8675309" } @@ -325,10 +326,25 @@ def test_drain_check_draining(self): self.runner.check_drained) def test_drain_check(self): + # Include a "metrics" instance to show that we ignore it. + self.plan["metrics_options"] = dict(enabled=True) + self.mock_ecs.locate_metrics_service.return_value = { + "deployments": [{ + "desiredCount": 1, + "runningCount": 1 + }], + "serviceArn": "Metric-Arn-01234567890" + } + mock_client = mock.Mock() mock_client.list_container_instances.side_effect = [ - {}, - {} + { # Actives + 'containerInstanceArns': [ + 'Metric-Arn-01234567890', + ], + "nextToken": "token-8675309" + }, + {} # Draining ] self.mock_boto.client.return_value = mock_client self.runner.check_drained()