From 91cf91f32d6462343b8130c6e72af5389b8ae351 Mon Sep 17 00:00:00 2001
From: Adrien Fillon <adrien.fillon@manomano.com>
Date: Mon, 19 Jul 2021 14:47:03 +0200
Subject: [PATCH] Fix deployment rollbacking with more than 100 tasks

The `list_tasks` call was capped at 100 tasks, if your service had a desired
superior of 100 the deployment would fail.

Signed-off-by: Adrien Fillon <adrien.fillon@manomano.com>
---
 ecs_deploy/ecs.py | 16 ++++++++++------
 tests/test_ecs.py | 14 ++++----------
 2 files changed, 14 insertions(+), 16 deletions(-)

diff --git a/ecs_deploy/ecs.py b/ecs_deploy/ecs.py
index cee8b62..3b9124d 100644
--- a/ecs_deploy/ecs.py
+++ b/ecs_deploy/ecs.py
@@ -3,6 +3,7 @@
 import re
 import copy
 from collections import defaultdict
+import itertools
 import logging
 import click_log
 
@@ -71,10 +72,13 @@ def describe_task_definition(self, task_definition_arn):
             )
 
     def list_tasks(self, cluster_name, service_name):
-        return self.boto.list_tasks(
-            cluster=cluster_name,
-            serviceName=service_name
-        )
+        tasks_paginator = self.boto.get_paginator(u'list_tasks')
+        return list(itertools.chain.from_iterable(
+            res['taskArns'] for res in tasks_paginator.paginate(
+                cluster=cluster_name,
+                serviceName=service_name
+            )
+        ))
 
     def describe_tasks(self, cluster_name, task_arns):
         return self.boto.describe_tasks(cluster=cluster_name, tasks=task_arns)
@@ -1173,11 +1177,11 @@ def is_deployed(self, service):
             cluster_name=service.cluster,
             service_name=service.name
         )
-        if not running_tasks[u'taskArns']:
+        if not running_tasks:
             return service.desired_count == 0
         running_count = self.get_running_tasks_count(
             service=service,
-            task_arns=running_tasks[u'taskArns']
+            task_arns=running_tasks
         )
         return service.desired_count == running_count
 
diff --git a/tests/test_ecs.py b/tests/test_ecs.py
index 90a2a5e..c5e4b3b 100644
--- a/tests/test_ecs.py
+++ b/tests/test_ecs.py
@@ -386,17 +386,11 @@
     u'test-task': RESPONSE_TASK_DEFINITION_2,
 }
 
-RESPONSE_LIST_TASKS_2 = {
-    u"taskArns": [TASK_ARN_1, TASK_ARN_2]
-}
+RESPONSE_LIST_TASKS_2 = [TASK_ARN_1, TASK_ARN_2]
 
-RESPONSE_LIST_TASKS_1 = {
-    u"taskArns": [TASK_ARN_1]
-}
+RESPONSE_LIST_TASKS_1 = [TASK_ARN_1]
 
-RESPONSE_LIST_TASKS_0 = {
-    u"taskArns": []
-}
+RESPONSE_LIST_TASKS_0 = []
 
 RESPONSE_DESCRIBE_TASKS = {
     u"tasks": [PAYLOAD_TASK_1, PAYLOAD_TASK_2]
@@ -1001,7 +995,7 @@ def test_client_describe_unknown_task_definition(client):
 
 def test_client_list_tasks(client):
     client.list_tasks(u'test-cluster', u'test-service')
-    client.boto.list_tasks.assert_called_once_with(cluster=u'test-cluster', serviceName=u'test-service')
+    client.boto.get_paginator.assert_called_once_with(u'list_tasks')
 
 
 def test_client_describe_tasks(client):