From 3f20ad2c13acbf1dce8f7decacb3a7a6d30e9db6 Mon Sep 17 00:00:00 2001
From: William Richard <william.richard.no.s@gmail.com>
Date: Tue, 20 Jun 2017 16:22:34 -0400
Subject: [PATCH] Support filtering by image id or image tag when describing
 ecr images

---
 moto/ecr/models.py               | 19 ++++++++---
 moto/ecr/responses.py            |  3 +-
 tests/test_ecr/test_ecr_boto3.py | 54 ++++++++++++++++++++++++++++++++
 3 files changed, 71 insertions(+), 5 deletions(-)

diff --git a/moto/ecr/models.py b/moto/ecr/models.py
index cbe8b2565..b90700ff4 100644
--- a/moto/ecr/models.py
+++ b/moto/ecr/models.py
@@ -193,16 +193,27 @@ class ECRBackend(BaseBackend):
                 images.append(image)
         return images
 
-    def describe_images(self, repository_name, registry_id=None, image_id=None):
+    def describe_images(self, repository_name, registry_id=None, image_ids=None):
 
         if repository_name in self.repositories:
             repository = self.repositories[repository_name]
         else:
             raise Exception("{0} is not a repository".format(repository_name))
 
-        response = []
-        for image in repository.images:
-            response.append(image)
+        if image_ids:
+            response = set()
+            for image_id in image_ids:
+                if 'imageDigest' in image_id:
+                    desired_digest = image_id['imageDigest']
+                    response.update([i for i in repository.images if i.get_image_digest() == desired_digest])
+                if 'imageTag' in image_id:
+                    desired_tag = image_id['imageTag']
+                    response.update([i for i in repository.images if i.image_tag == desired_tag])
+        else:
+            response = []
+            for image in repository.images:
+                response.append(image)
+
         return response
 
     def put_image(self, repository_name, image_manifest, image_tag):
diff --git a/moto/ecr/responses.py b/moto/ecr/responses.py
index f8b1606cc..40d8cfb66 100644
--- a/moto/ecr/responses.py
+++ b/moto/ecr/responses.py
@@ -69,7 +69,8 @@ class ECRResponse(BaseResponse):
     def describe_images(self):
         repository_str = self._get_param('repositoryName')
         registry_id = self._get_param('registryId')
-        images = self.ecr_backend.describe_images(repository_str, registry_id)
+        image_ids = self._get_param('imageIds')
+        images = self.ecr_backend.describe_images(repository_str, registry_id, image_ids)
         return json.dumps({
             'imageDetails': [image.response_describe_object for image in images],
         })
diff --git a/tests/test_ecr/test_ecr_boto3.py b/tests/test_ecr/test_ecr_boto3.py
index 3a32c1515..647015446 100644
--- a/tests/test_ecr/test_ecr_boto3.py
+++ b/tests/test_ecr/test_ecr_boto3.py
@@ -314,3 +314,57 @@ def test_describe_images():
     response['imageDetails'][0]['imageSizeInBytes'].should.equal(52428800)
     response['imageDetails'][1]['imageSizeInBytes'].should.equal(52428800)
     response['imageDetails'][2]['imageSizeInBytes'].should.equal(52428800)
+
+
+@mock_ecr
+def test_describe_images_by_tag():
+    client = boto3.client('ecr', region_name='us-east-1')
+    _ = client.create_repository(
+        repositoryName='test_repository'
+    )
+
+    tag_map = {}
+    for tag in ['latest', 'v1', 'v2']:
+        put_response = client.put_image(
+            repositoryName='test_repository',
+            imageManifest=json.dumps(_create_image_manifest()),
+            imageTag=tag
+        )
+        tag_map[tag] = put_response['image']
+
+    for tag, put_response in tag_map.items():
+        response = client.describe_images(repositoryName='test_repository', imageIds=[{'imageTag': tag}])
+        len(response['imageDetails']).should.be(1)
+        image_detail = response['imageDetails'][0]
+        image_detail['registryId'].should.equal("012345678910")
+        image_detail['repositoryName'].should.equal("test_repository")
+        image_detail['imageTags'].should.equal([put_response['imageId']['imageTag']])
+        image_detail['imageDigest'].should.equal(put_response['imageId']['imageDigest'])
+
+
+@mock_ecr
+def test_describe_images_by_digest():
+    client = boto3.client('ecr', region_name='us-east-1')
+    _ = client.create_repository(
+        repositoryName='test_repository'
+    )
+
+    tags = ['latest', 'v1', 'v2']
+    digest_map = {}
+    for tag in tags:
+        put_response = client.put_image(
+            repositoryName='test_repository',
+            imageManifest=json.dumps(_create_image_manifest()),
+            imageTag=tag
+        )
+        digest_map[put_response['image']['imageId']['imageDigest']] = put_response['image']
+
+    for digest, put_response in digest_map.items():
+        response = client.describe_images(repositoryName='test_repository',
+                                          imageIds=[{'imageDigest': digest}])
+        len(response['imageDetails']).should.be(1)
+        image_detail = response['imageDetails'][0]
+        image_detail['registryId'].should.equal("012345678910")
+        image_detail['repositoryName'].should.equal("test_repository")
+        image_detail['imageTags'].should.equal([put_response['imageId']['imageTag']])
+        image_detail['imageDigest'].should.equal(digest)