From 1697706ac0f536e6905d6a724f0fa5e17a681485 Mon Sep 17 00:00:00 2001
From: zhuwenxing <wenxing.zhu@zilliz.com>
Date: Wed, 12 Jun 2024 14:27:56 +0800
Subject: [PATCH] test: [skip e2e] optimize get_image_tag_by_short_name
 function (#33780)

Refactor the function to improve performance and readability. Instead of
making API requests to Docker Hub, the function now retrieves tags from
the Harbor registry. It also filters the tags based on the provided
architecture and selects the latest tag that matches the prefix. This
change enhances the efficiency of retrieving image tags by short name.

Signed-off-by: zhuwenxing <wenxing.zhu@zilliz.com>
---
 tests/scripts/get_image_tag_by_short_name.py | 54 ++++++++------------
 1 file changed, 22 insertions(+), 32 deletions(-)

diff --git a/tests/scripts/get_image_tag_by_short_name.py b/tests/scripts/get_image_tag_by_short_name.py
index 10ca09e8100f5..d3701e4369661 100644
--- a/tests/scripts/get_image_tag_by_short_name.py
+++ b/tests/scripts/get_image_tag_by_short_name.py
@@ -3,44 +3,34 @@
 from tenacity import retry, stop_after_attempt
 
 @retry(stop=stop_after_attempt(7))
-def get_image_tag_by_short_name(repository, tag, arch):
+def get_image_tag_by_short_name(tag, arch):
 
-    # Send API request to get all tags start with prefix
-    # ${branch}-latest  means the tag is a dev build
-        # master-latest -> master-$date-$commit
-        # 2.3.0-latest -> 2.3.0-$date-$commit
-    # latest means the tag is a release build
-        # latest -> v$version
-    splits = tag.split("-")
-    prefix = f"{splits[0]}-" if len(splits) > 1 else "v"
-    url = f"https://hub.docker.com/v2/repositories/{repository}/tags?name={prefix}&ordering=last_updated"
-    response = requests.get(url)
-    data = response.json()
-
-    # Get the latest tag with the same arch and prefix
-    sorted_images = sorted(data["results"], key=lambda x: x["last_updated"], reverse=True)
-    # print(sorted_images)
-    candidate_tag = None
-    for tag_info in sorted_images:
-        # print(tag_info)
-        if tag == "2.2.0-latest": # special case for 2.2.0-latest, for 2.2.0 branch, there is no arm amd and gpu as suffix
-            candidate_tag = tag_info["name"]
-        else:
-            if arch in tag_info["name"]:
-                candidate_tag = tag_info["name"]
-            else:
-                continue
-        if candidate_tag == tag:
-            continue
-        else:
-            return candidate_tag
+    prefix = tag.split("-")[0]
+    url = f"https://harbor.milvus.io/api/v2.0/projects/milvus/repositories/milvus/artifacts?with_tag=true&q=tags%253D~{prefix}-&page_size=100&page=1"
 
+    payload = {}
+    response = requests.request("GET", url, data=payload)
+    rsp = response.json()
+    tag_list = []
+    for r in rsp:
+        tags = r["tags"]
+        for tag in tags:
+            tag_list.append(tag["name"])
+    tag_candidates = []
+    for t in tag_list:
+        r = t.split("-")
+        if len(r) == 4 and arch in t:
+            tag_candidates.append(t)
+    tag_candidates.sort()
+    if len(tag_candidates) == 0:
+        return tag
+    else:
+        return tag_candidates[-1]
 
 if __name__ == "__main__":
     argparse = argparse.ArgumentParser()
-    argparse.add_argument("--repository", type=str, default="milvusdb/milvus")
     argparse.add_argument("--tag", type=str, default="master-latest")
     argparse.add_argument("--arch", type=str, default="amd64")
     args = argparse.parse_args()
-    res = get_image_tag_by_short_name(args.repository, args.tag, args.arch)
+    res = get_image_tag_by_short_name(args.tag, args.arch)
     print(res)