From 7c1fd0a2f1b834dcc7cd6e9c852cff765b066cc7 Mon Sep 17 00:00:00 2001
From: Steve Pulec <spulec@gmail.com>
Date: Sun, 24 Jun 2018 20:13:39 -0400
Subject: [PATCH] Fix ECS update_service and describing tasks.

---
 moto/ecs/exceptions.py           | 11 +++++++++++
 moto/ecs/models.py               | 10 ++++++----
 tests/test_ecs/test_ecs_boto3.py | 23 +++++++++++++++++++++++
 3 files changed, 40 insertions(+), 4 deletions(-)
 create mode 100644 moto/ecs/exceptions.py

diff --git a/moto/ecs/exceptions.py b/moto/ecs/exceptions.py
new file mode 100644
index 000000000..c23d6fd1d
--- /dev/null
+++ b/moto/ecs/exceptions.py
@@ -0,0 +1,11 @@
+from __future__ import unicode_literals
+from moto.core.exceptions import RESTError
+
+
+class ServiceNotFoundException(RESTError):
+    code = 400
+
+    def __init__(self, service_name):
+        super(ServiceNotFoundException, self).__init__(
+            error_type="ServiceNotFoundException",
+            message="The service {0} does not exist".format(service_name))
diff --git a/moto/ecs/models.py b/moto/ecs/models.py
index 3c51cd03f..b8dc1b738 100644
--- a/moto/ecs/models.py
+++ b/moto/ecs/models.py
@@ -10,6 +10,8 @@ from moto.core import BaseBackend, BaseModel
 from moto.ec2 import ec2_backends
 from copy import copy
 
+from .exceptions import ServiceNotFoundException
+
 
 class BaseObject(BaseModel):
 
@@ -601,8 +603,9 @@ class EC2ContainerServiceBackend(BaseBackend):
             raise Exception("tasks cannot be empty")
         response = []
         for cluster, cluster_tasks in self.tasks.items():
-            for task_id, task in cluster_tasks.items():
-                if task_id in tasks or task.task_arn in tasks:
+            for task_arn, task in cluster_tasks.items():
+                task_id = task_arn.split("/")[-1]
+                if task_arn in tasks or task.task_arn in tasks or any(task_id in task for task in tasks):
                     response.append(task)
         return response
 
@@ -698,8 +701,7 @@ class EC2ContainerServiceBackend(BaseBackend):
                     cluster_service_pair].desired_count = desired_count
             return self.services[cluster_service_pair]
         else:
-            raise Exception("cluster {0} or service {1} does not exist".format(
-                cluster_name, service_name))
+            raise ServiceNotFoundException(service_name)
 
     def delete_service(self, cluster_name, service_name):
         cluster_service_pair = '{0}:{1}'.format(cluster_name, service_name)
diff --git a/tests/test_ecs/test_ecs_boto3.py b/tests/test_ecs/test_ecs_boto3.py
index d2cfd3724..b2d06d035 100644
--- a/tests/test_ecs/test_ecs_boto3.py
+++ b/tests/test_ecs/test_ecs_boto3.py
@@ -2,6 +2,7 @@ from __future__ import unicode_literals
 
 from copy import deepcopy
 
+from botocore.exceptions import ClientError
 import boto3
 import sure  # noqa
 import json
@@ -450,6 +451,21 @@ def test_update_service():
     response['service']['desiredCount'].should.equal(0)
 
 
+@mock_ecs
+def test_update_missing_service():
+    client = boto3.client('ecs', region_name='us-east-1')
+    _ = client.create_cluster(
+        clusterName='test_ecs_cluster'
+    )
+
+    client.update_service.when.called_with(
+        cluster='test_ecs_cluster',
+        service='test_ecs_service',
+        taskDefinition='test_ecs_task',
+        desiredCount=0
+    ).should.throw(ClientError)
+
+
 @mock_ecs
 def test_delete_service():
     client = boto3.client('ecs', region_name='us-east-1')
@@ -1054,6 +1070,13 @@ def test_describe_tasks():
     set([response['tasks'][0]['taskArn'], response['tasks']
     [1]['taskArn']]).should.equal(set(tasks_arns))
 
+    # Test we can pass task ids instead of ARNs
+    response = client.describe_tasks(
+        cluster='test_ecs_cluster',
+        tasks=[tasks_arns[0].split("/")[-1]]
+    )
+    len(response['tasks']).should.equal(1)
+
 
 @mock_ecs
 def describe_task_definition():