From 8029a46f253c8863d1cf43d6d57083c5fcd8b5ed Mon Sep 17 00:00:00 2001
From: jrconlin <jconlin+git@mozilla.com>
Date: Mon, 10 Apr 2017 16:13:42 -0700
Subject: [PATCH] bug: Ensure all containers drained before exiting

(wip)

closes #45
---
 ardere/exceptions.py         |  4 ++++
 ardere/step_functions.py     | 28 ++++++++++++++++++++++++++
 handler.py                   |  4 ++++
 serverless.yml               | 18 +++++++++++++++++
 tests/test_step_functions.py | 39 ++++++++++++++++++++++++++++++++++++
 5 files changed, 93 insertions(+)

diff --git a/ardere/exceptions.py b/ardere/exceptions.py
index af5cfc6..1cd0288 100644
--- a/ardere/exceptions.py
+++ b/ardere/exceptions.py
@@ -8,3 +8,7 @@ class ShutdownPlanException(Exception):
 
 class ValidationException(Exception):
     """Exception to indicate validation error parsing input"""
+
+
+class UndrainedInstancesException(Exception):
+    """There are still ACTIVE or DRAINING instances in the cluster"""
diff --git a/ardere/step_functions.py b/ardere/step_functions.py
index f2c3d59..1100f50 100644
--- a/ardere/step_functions.py
+++ b/ardere/step_functions.py
@@ -26,6 +26,7 @@
     ServicesStartingException,
     ShutdownPlanException,
     ValidationException,
+    UndrainedInstancesException,
 )
 
 logger = logging.getLogger()
@@ -385,3 +386,30 @@ def cleanup_cluster(self):
         except botocore.exceptions.ClientError:
             pass
         return self.event
+
+    def check_drained(self):
+        """Ensure that all services are shut down before allowing restart
+
+        Step 8
+
+        """
+        client = self.boto.client('ecs')
+        actives = len(
+            client.list_container_instances(
+                cluster=self.event["ecs_name"],
+                maxResults=1,
+                status="ACTIVE",
+            ).get('containerInstanceArns', []))
+        if actives:
+            raise UndrainedInstancesException(
+                "Still {} active.".format(actives))
+        draining = len(
+            client.list_container_instances(
+                cluster=self.event["ecs_name"],
+                maxResults=1,
+                status="DRAINING",
+            ).get('containerInstanceArns', []))
+        if draining:
+            raise UndrainedInstancesException(
+                "Still {} draining.".format(draining))
+        return self.event
diff --git a/handler.py b/handler.py
index 245726e..1ad268d 100644
--- a/handler.py
+++ b/handler.py
@@ -39,3 +39,7 @@ def check_for_cluster_done(event, context):
 def cleanup_cluster(event, context):
     runner = AsynchronousPlanRunner(event, context)
     return runner.cleanup_cluster()
+
+
+def check_drain(event, context):
+    return AsynchronousPlanRunner(event, context).check_drained()
\ No newline at end of file
diff --git a/serverless.yml b/serverless.yml
index 397eae3..d2e6034 100644
--- a/serverless.yml
+++ b/serverless.yml
@@ -107,6 +107,8 @@ functions:
     handler: handler.check_for_cluster_done
   cleanup_cluster:
     handler: handler.cleanup_cluster
+  check_drain:
+    handler: handler.check_drain
 
 stepFunctions:
   stateMachines:
@@ -196,7 +198,23 @@ stepFunctions:
         "Clean-up Cluster":
           Type: Task
           Resource: cleanup_cluster
+          Next: "Checking Drain"
+        "Checking Drain":
+          Type: Task
+          Resource: check_drain
           End: true
+          Retry:
+            -
+              ErrorEquals:
+                - UndrainedInstancesException
+              IntervalSeconds: 10
+              MaxAttempts: 10
+              BackoffRate: 1
+          Catch:
+            -
+              ErrorEquals:
+                - States.ALL
+              ResultPath: "$.error-info"
 
 resources:
   Resources:
diff --git a/tests/test_step_functions.py b/tests/test_step_functions.py
index f59fc2b..c71ee72 100644
--- a/tests/test_step_functions.py
+++ b/tests/test_step_functions.py
@@ -273,6 +273,45 @@ def test_cleanup_cluster_error(self):
         self.runner.cleanup_cluster()
         mock_s3.Object.assert_called()
 
+    def test_drain_check_active(self):
+        from ardere.exceptions import UndrainedInstancesException
+
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.return_value = {
+            'containerInstanceArns': [
+                'Some-Arn-01234567890',
+            ],
+            "nextToken": "token-8675309"
+        }
+        self.mock_boto.client.return_value = mock_client
+        assert_raises(UndrainedInstancesException,
+                      self.runner.check_drained)
+
+    def test_drain_check_draining(self):
+        from ardere.exceptions import UndrainedInstancesException
+
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.side_effect = [
+            {},
+            {
+                'containerInstanceArns': [
+                    'Some-Arn-01234567890',
+                ],
+                "nextToken": "token-8675309"
+            }
+        ]
+        self.mock_boto.client.return_value = mock_client
+        assert_raises(UndrainedInstancesException,
+                      self.runner.check_drained)
+
+    def test_drain_check(self):
+        mock_client = mock.Mock()
+        mock_client.list_container_instances.side_effect = [
+            {},
+            {}
+        ]
+        self.mock_boto.client.return_value = mock_client
+        self.runner.check_drained()
 
 class TestValidation(unittest.TestCase):
     def _make_FUT(self):