Internal change

PiperOrigin-RevId: 374922602
Change-Id: I3ac804f7436775a14d4905442c36201e3fc7dce0

5 files changed, 150 insertions, 18 deletions

diff --git a/tradefed_cluster/BUILD b/tradefed_cluster/BUILD
index 8faa426..cc9a5c2 100644
--- a/tradefed_cluster/BUILD
+++ b/tradefed_cluster/BUILD
@@ -268,6 +268,7 @@ pytype_strict_library(
         ":datastore_entities_lib",
         ":datastore_util_lib",
         ":device_manager_lib",
+        ":harness_image_metadata_syncer_lib",
         ":host_event_lib",
         ":note_manager_lib",
         "//tradefed_cluster/util:ndb_shim_lib",
@@ -1243,8 +1244,10 @@ py_test(
     srcs_version = "PY3",
     deps = [
         ":datastore_entities_lib",
+        ":datastore_test_util_lib",
         ":harness_image_metadata_syncer_lib",
         ":testbed_dependent_test_lib",
+        "//third_party/py/absl/testing:parameterized",
         "//tradefed_cluster/util:ndb_shim_lib",
         third_party("mock"),
     ],
diff --git a/tradefed_cluster/cluster_host_api.py b/tradefed_cluster/cluster_host_api.py
index 6d1f4f7..8ef21d6 100644
--- a/tradefed_cluster/cluster_host_api.py
+++ b/tradefed_cluster/cluster_host_api.py
@@ -28,6 +28,7 @@ from tradefed_cluster import common
 from tradefed_cluster import datastore_entities
 from tradefed_cluster import datastore_util
 from tradefed_cluster import device_manager
+from tradefed_cluster import harness_image_metadata_syncer
 from tradefed_cluster import host_event
 from tradefed_cluster import note_manager
 
@@ -904,7 +905,8 @@ class ClusterHostApi(remote.Service):
       if not metadata:
         metadata = datastore_entities.HostMetadata(
             id=hostname, hostname=hostname)
-      if metadata.test_harness_image != request.test_harness_image:
+      if not harness_image_metadata_syncer.AreHarnessImagesEqual(
+          metadata.test_harness_image, request.test_harness_image):
         event = host_event.HostEvent(
             time=datetime.datetime.utcnow(),
             type=_HOST_UPDATE_STATE_CHANGED_EVENT_NAME,
@@ -927,4 +929,3 @@ class ClusterHostApi(remote.Service):
       error_message += ("Hosts [%s] are not enabled to be updated from UI. "
                         % ", ".join(hosts_not_enabled))
     raise endpoints.BadRequestException(error_message)
-
diff --git a/tradefed_cluster/cluster_host_api_test.py b/tradefed_cluster/cluster_host_api_test.py
index 30baa52..d7be3d0 100644
--- a/tradefed_cluster/cluster_host_api_test.py
+++ b/tradefed_cluster/cluster_host_api_test.py
@@ -2148,14 +2148,19 @@ class ClusterHostApiTest(api_test.ApiTest):
     hostname2 = 'host2'
     lab_name = 'alab'
     owner = 'user1'
-    image_name = 'image_url'
+    repo_name = 'test_repo'
+    tag = 'new'
+    image_name_new = ':'.join([repo_name, tag])
     datastore_test_util.CreateHostConfig(
         hostname1, lab_name, owners=[owner], enable_ui_update=True)
     datastore_test_util.CreateHostConfig(
         hostname2, lab_name, owners=[owner], enable_ui_update=True)
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d2', current_tags=[tag])
+
     api_request = {
         'hostnames': [hostname1, hostname2],
-        'test_harness_image': image_name,
+        'test_harness_image': image_name_new,
         'user': owner,
     }
     api_response = self.testapp.post_json(
@@ -2165,14 +2170,14 @@ class ClusterHostApiTest(api_test.ApiTest):
 
     host_metadata = datastore_entities.HostMetadata.get_by_id(hostname1)
     self.assertEqual(hostname1, host_metadata.hostname)
-    self.assertEqual(image_name, host_metadata.test_harness_image)
+    self.assertEqual(image_name_new, host_metadata.test_harness_image)
     host_update_state = datastore_entities.HostUpdateState.get_by_id(hostname1)
     self.assertEqual(hostname1, host_update_state.hostname)
     self.assertEqual(api_messages.HostUpdateState.PENDING,
                      host_update_state.state)
     host_metadata = datastore_entities.HostMetadata.get_by_id(hostname2)
     self.assertEqual(hostname2, host_metadata.hostname)
-    self.assertEqual(image_name, host_metadata.test_harness_image)
+    self.assertEqual(image_name_new, host_metadata.test_harness_image)
     host_update_state = datastore_entities.HostUpdateState.get_by_id(hostname2)
     self.assertEqual(hostname2, host_update_state.hostname)
     self.assertEqual(api_messages.HostUpdateState.PENDING,
@@ -2184,8 +2189,11 @@ class ClusterHostApiTest(api_test.ApiTest):
     hostname2 = 'host2'
     lab_name = 'alab'
     owner = 'user1'
-    image_name_old = 'image_url_old'
-    image_name_new = 'image_url_new'
+    repo_name = 'test_repo'
+    tag_old = 'old'
+    tag_new = 'new'
+    image_name_old = ':'.join([repo_name, tag_old])
+    image_name_new = ':'.join([repo_name, tag_new])
     datastore_test_util.CreateHostConfig(
         hostname1, lab_name, owners=[owner], enable_ui_update=True)
     datastore_test_util.CreateHostConfig(
@@ -2194,6 +2202,11 @@ class ClusterHostApiTest(api_test.ApiTest):
         hostname=hostname1, test_harness_image=image_name_old)
     datastore_test_util.CreateHostMetadata(
         hostname=hostname2, test_harness_image=image_name_old)
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d1', current_tags=[tag_old])
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d2', current_tags=[tag_new])
+
     api_request = {
         'hostnames': [hostname1, hostname2],
         'test_harness_image': image_name_new,
@@ -2219,19 +2232,26 @@ class ClusterHostApiTest(api_test.ApiTest):
     self.assertEqual(api_messages.HostUpdateState.PENDING,
                      host_update_state.state)
 
-  def testBatchUpdateHostMetadata_succeedsWithUnchangedMetadata(self):
+  def testBatchUpdateHostMetadata_succeedsWithUnchangedImage(self):
     """Test batch set test_harness_image for hosts."""
     hostname = 'host1'
     lab_name = 'alab'
     owner = 'user1'
-    image_name = 'image_url'
+    repo_name = 'test_repo'
+    tag_old = 'old'
+    tag_new = 'new'
+    image_name_old = ':'.join([repo_name, tag_old])
+    image_name_new = ':'.join([repo_name, tag_new])
     datastore_test_util.CreateHostConfig(
         hostname, lab_name, owners=[owner], enable_ui_update=True)
     datastore_test_util.CreateHostMetadata(
-        hostname=hostname, test_harness_image=image_name)
+        hostname=hostname, test_harness_image=image_name_old)
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d1', current_tags=[tag_old, tag_new])
+
     api_request = {
         'hostnames': [hostname],
-        'test_harness_image': image_name,
+        'test_harness_image': image_name_new,
         'user': owner,
     }
     api_response = self.testapp.post_json(
@@ -2241,7 +2261,7 @@ class ClusterHostApiTest(api_test.ApiTest):
 
     host_metadata = datastore_entities.HostMetadata.get_by_id(hostname)
     self.assertEqual(hostname, host_metadata.hostname)
-    self.assertEqual(image_name, host_metadata.test_harness_image)
+    self.assertEqual(image_name_new, host_metadata.test_harness_image)
     # No update should start in this case, because the image is not changed.
     self.assertIsNone(datastore_entities.HostUpdateState.get_by_id(hostname))
 
@@ -2251,8 +2271,11 @@ class ClusterHostApiTest(api_test.ApiTest):
     hostname2 = 'host2'
     lab_name = 'alab'
     owner = 'user1'
-    image_name_old = 'image_url_old'
-    image_name_new = 'image_url_new'
+    repo_name = 'test_repo'
+    tag_old = 'old'
+    tag_new = 'new'
+    image_name_old = ':'.join([repo_name, tag_old])
+    image_name_new = ':'.join([repo_name, tag_new])
     datastore_test_util.CreateHostConfig(
         hostname1, lab_name, owners=[owner], enable_ui_update=True)
     datastore_test_util.CreateHostConfig(
@@ -2261,6 +2284,11 @@ class ClusterHostApiTest(api_test.ApiTest):
         hostname=hostname1, test_harness_image=image_name_old)
     datastore_test_util.CreateHostMetadata(
         hostname=hostname2, test_harness_image=image_name_old)
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d1', current_tags=[tag_old])
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d2', current_tags=[tag_new])
+
     api_request = {
         'hostnames': [hostname1, hostname2],
         'test_harness_image': image_name_new,
@@ -2283,8 +2311,11 @@ class ClusterHostApiTest(api_test.ApiTest):
     hostname2 = 'host2'
     lab_name = 'alab'
     owner = 'user1'
-    image_name_old = 'image_url_old'
-    image_name_new = 'image_url_new'
+    repo_name = 'test_repo'
+    tag_old = 'old'
+    tag_new = 'new'
+    image_name_old = ':'.join([repo_name, tag_old])
+    image_name_new = ':'.join([repo_name, tag_new])
     datastore_test_util.CreateHostConfig(
         hostname1, lab_name, owners=[owner], enable_ui_update=False)
     datastore_test_util.CreateHostConfig(
@@ -2293,6 +2324,11 @@ class ClusterHostApiTest(api_test.ApiTest):
         hostname=hostname1, test_harness_image=image_name_old)
     datastore_test_util.CreateHostMetadata(
         hostname=hostname2, test_harness_image=image_name_old)
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d1', current_tags=[tag_old])
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name=repo_name, digest='d2', current_tags=[tag_new])
+
     api_request = {
         'hostnames': [hostname1, hostname2],
         'test_harness_image': image_name_new,
diff --git a/tradefed_cluster/harness_image_metadata_syncer.py b/tradefed_cluster/harness_image_metadata_syncer.py
index a47f9ee..45f861d 100644
--- a/tradefed_cluster/harness_image_metadata_syncer.py
+++ b/tradefed_cluster/harness_image_metadata_syncer.py
@@ -121,3 +121,68 @@ def _AnalyseTradefedDockerImageTags(tags):
     if _HISTORICAL_GOLDEN_PATTERN.match(tag):
       result['is_historical_golden'] = True
   return result
+
+
+def _GetTestHarnessImageMetadata(repo_name, image_tag):
+  """Get a TestHarnessImageMetadata entity.
+
+  Args:
+    repo_name: string, repo name.
+    image_tag: string, image tag.
+
+  Returns:
+    A entity of TestHarnessImageMetadata, or None.
+  """
+  result = None
+  response = list(datastore_entities.TestHarnessImageMetadata.query().filter(
+      datastore_entities.TestHarnessImageMetadata.repo_name ==
+      repo_name).filter(datastore_entities.TestHarnessImageMetadata.current_tags
+                        == image_tag).fetch(1))
+  if response:
+    result = response[0]
+  return result
+
+
+def _SplitImageUrlIntoRepoAndTag(image_url):
+  """Split image url to tuple of repo_name and image_tag.
+
+  Args:
+    image_url: string, image url.
+
+  Returns:
+    string, repo name.
+    string, image tag name.
+  """
+  delimiter = ':'
+  default_tag = 'latest'
+  if delimiter in image_url:
+    repo_name, image_tag = image_url.split(delimiter, 1)
+  else:
+    repo_name = image_url
+    image_tag = default_tag
+  return repo_name, image_tag
+
+
+def AreHarnessImagesEqual(image_url_a, image_url_b):
+  """Helper function to check whether the images behind urls equal.
+
+  Args:
+    image_url_a: string, url to image A.
+    image_url_b: string, url to image B.
+
+  Returns:
+    bool, whether two images urls links to the same image digest.
+  """
+
+  if image_url_a == image_url_b:
+    return True
+  if not image_url_a or not image_url_b:
+    return False
+
+  repo_a, tag_a = _SplitImageUrlIntoRepoAndTag(image_url_a)
+  repo_b, tag_b = _SplitImageUrlIntoRepoAndTag(image_url_b)
+  image_a = _GetTestHarnessImageMetadata(repo_a, tag_a)
+  image_b = _GetTestHarnessImageMetadata(repo_b, tag_b)
+  if image_a and image_b and image_a.digest == image_b.digest:
+    return True
+  return False
diff --git a/tradefed_cluster/harness_image_metadata_syncer_test.py b/tradefed_cluster/harness_image_metadata_syncer_test.py
index aad40bd..0b2c342 100644
--- a/tradefed_cluster/harness_image_metadata_syncer_test.py
+++ b/tradefed_cluster/harness_image_metadata_syncer_test.py
@@ -14,15 +14,20 @@
 """Tests for harness_image_metadata_syncer."""
 import datetime
 import unittest
+
+from absl.testing import parameterized
+
 import mock
 
 from tradefed_cluster import datastore_entities
+from tradefed_cluster import datastore_test_util
 from tradefed_cluster import harness_image_metadata_syncer
 from tradefed_cluster import testbed_dependent_test
 from tradefed_cluster.util import ndb_shim as ndb
 
 
-class HarnessImageMetadataSyncerTest(testbed_dependent_test.TestbedDependentTest
+class HarnessImageMetadataSyncerTest(parameterized.TestCase,
+                                     testbed_dependent_test.TestbedDependentTest
                                     ):
   """Tests for HarnessImageMetadataSyncer."""
 
@@ -213,6 +218,28 @@ class HarnessImageMetadataSyncerTest(testbed_dependent_test.TestbedDependentTest
     self.assertCountEqual(['3333333', 'staging'], entity_3.current_tags)
     self.assertEmpty(entity_3.historical_tags)
 
+  @parameterized.named_parameters(
+      ('Same image digest.', 'test_repo:t1', 'test_repo:t2', True),
+      ('Different image digests.', 'test_repo:t1', 'test_repo:t3', False),
+      ('Only split on first delimiter.', 'test_repo:t1:xx', 'test_repo:t1',
+       False),
+      ('Same url', 'test_repo:t1', 'test_repo:t1', True),
+      ('Same url with default tag', 'test_repo:t3', 'test_repo', True),
+      ('One is empty', '', 'test_repo:t1', False),
+      ('Image not found', 'test_repo:notfound', 'test_repo:t1', False),
+  )
+  def testAreHarnessImagesEqual(self, image_url_a, image_url_b,
+                                expected_result):
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name='test_repo', digest='d1', current_tags=['t1', 't2'])
+    datastore_test_util.CreateTestHarnessImageMetadata(
+        repo_name='test_repo', digest='d2', current_tags=['t3', 'latest'])
+
+    self.assertEqual(
+        expected_result,
+        harness_image_metadata_syncer.AreHarnessImagesEqual(
+            image_url_a, image_url_b))
+
 
 if __name__ == '__main__':
   unittest.main()