Add a batch get_matching_tags_for_images method

This will be used in the security notification worker to retrieving the tags needed in a set of batch calls, rather than multiple calls per image

data/model/test/test_tag.py

@@ -1,25 +1,84 @@
-from data.database import Image, RepositoryTag
+from data.database import Image, RepositoryTag, ImageStorage, Repository
 from data.model.repository import create_repository
 from data.model.tag import (list_active_repo_tags, create_or_update_tag, delete_tag,
-                            get_matching_tags, _tag_alive)
+                            get_matching_tags, _tag_alive, get_matching_tags_for_images)
 from data.model.image import find_create_or_link_image
 
 from test.fixtures import *
 
+def _get_expected_tags(image):
+  expected_query = (RepositoryTag
+                  .select()
+                  .join(Image)
+                  .where(RepositoryTag.hidden == False)
+                  .where((Image.id == image.id) | (Image.ancestors ** ('%%/%s/%%' % image.id))))
+  return set([tag.id for tag in _tag_alive(expected_query)])
+
+
 def test_get_matching_tags(initialized_db):
   # Test for every image in the test database.
-  for image in Image.select():
+  for image in Image.select(Image, ImageStorage).join(ImageStorage):
     matching_query = get_matching_tags(image.docker_image_id, image.storage.uuid)
-    expected_query = (RepositoryTag
-                      .select()
-                      .join(Image)
-                      .where(RepositoryTag.hidden == False)
-                      .where((Image.id == image.id) | (Image.ancestors ** ('%%/%s/%%' % image.id))))
-
     matching_tags = set([tag.id for tag in matching_query])
-    expected_tags = set([tag.id for tag in _tag_alive(expected_query)])
+    expected_tags = _get_expected_tags(image)
     assert matching_tags == expected_tags, "mismatch for image %s" % image.id
 
+
+def test_get_matching_tag_ids_for_images(initialized_db):
+  # Try for various sets of the first N images.
+  for count in [5, 10, 15]:
+    pairs = []
+    expected_tags_ids = set()
+    for image in Image.select(Image, ImageStorage).join(ImageStorage):
+      if len(pairs) >= count:
+        break
+
+      pairs.append((image.docker_image_id, image.storage.uuid))
+      expected_tags_ids.update(_get_expected_tags(image))
+
+    matching_tags_ids = set([tag.id for tag in get_matching_tags_for_images(pairs)])
+    assert matching_tags_ids == expected_tags_ids
+
+
+def test_get_matching_tag_ids_for_all_images(initialized_db):
+  pairs = []
+  for image in Image.select(Image, ImageStorage).join(ImageStorage):
+    pairs.append((image.docker_image_id, image.storage.uuid))
+
+  expected_tags_ids = set([tag.id for tag in _tag_alive(RepositoryTag.select())])
+  matching_tags_ids = set([tag.id for tag in get_matching_tags_for_images(pairs)])
+
+  # Ensure every alive tag was found.
+  assert matching_tags_ids == expected_tags_ids
+
+
+def test_get_matching_tag_ids_images_filtered(initialized_db):
+  def filter_query(query):
+    return query.join(Repository).where(Repository.name == 'simple')
+
+  filtered_images = filter_query(Image
+                                 .select(Image, ImageStorage)
+                                 .join(RepositoryTag)
+                                 .switch(Image)
+                                 .join(ImageStorage)
+                                 .switch(Image))
+
+  expected_tags_query = _tag_alive(filter_query(RepositoryTag
+                                                .select()))
+
+  pairs = []
+  for image in filtered_images:
+    pairs.append((image.docker_image_id, image.storage.uuid))
+
+  matching_tags = get_matching_tags_for_images(pairs, filter_query=filter_query)
+
+  expected_tag_ids = set([tag.id for tag in expected_tags_query])
+  matching_tags_ids = set([tag.id for tag in matching_tags])
+
+  # Ensure every alive tag was found.
+  assert matching_tags_ids == expected_tag_ids
+
+
 def assert_tags(repository, *args):
   tags = list(list_active_repo_tags(repository))
   assert len(tags) == len(args)