From 97525cc2fc887259e1a30e8a5f54cd948da7754c Mon Sep 17 00:00:00 2001
From: Cesar Gray Blanco <cegraybl@microsoft.com>
Date: Thu, 26 Sep 2024 14:20:13 -0700
Subject: [PATCH] fix the return values for 'list' command, allows output to be
 transformed by cli --output option

---
 src/acrcssc/azext_acrcssc/cssc.py             |  5 +--
 .../azext_acrcssc/helper/_workflow_status.py  | 45 ++++++++++++++-----
 2 files changed, 35 insertions(+), 15 deletions(-)

diff --git a/src/acrcssc/azext_acrcssc/cssc.py b/src/acrcssc/azext_acrcssc/cssc.py
index 5d5b6e60d00..ba227b76434 100644
--- a/src/acrcssc/azext_acrcssc/cssc.py
+++ b/src/acrcssc/azext_acrcssc/cssc.py
@@ -139,7 +139,6 @@ def list_scan_status(cmd, registry_name, resource_group_name, status, workflow_t
     registry = acr_client_registries.get(resource_group_name, registry_name)
 
     image_status = track_scan_progress(cmd, resource_group_name, registry, status)
-    for image in image_status:
-        print(image)
-
     print(f"Total images: {len(image_status)}")
+
+    return image_status
diff --git a/src/acrcssc/azext_acrcssc/helper/_workflow_status.py b/src/acrcssc/azext_acrcssc/helper/_workflow_status.py
index fd7f230d642..8af9b96275d 100644
--- a/src/acrcssc/azext_acrcssc/helper/_workflow_status.py
+++ b/src/acrcssc/azext_acrcssc/helper/_workflow_status.py
@@ -254,9 +254,10 @@ def from_taskrun(cmd, taskrun_client, registry, scan_taskruns, patch_taskruns, p
                 patch_task = next(task for task in patch_taskruns if task.run_id == patch_task_id)
                 all_status[image].patch_task = patch_task
 
-        return all_status.values()
+        # don't return a list of WorkflowTaskStatus object, 
+        return [status.get_status() for status in all_status.values()]
 
-    def __str__(self) -> str:
+    def get_status(self):
         scan_status = self.scan_status()
         scan_date = WORKFLOW_STATUS_NOT_AVAILABLE if self.scan_task is None else self.scan_task.create_time
         scan_task_id = WORKFLOW_STATUS_NOT_AVAILABLE if self.scan_task is None else self.scan_task.run_id
@@ -274,19 +275,39 @@ def __str__(self) -> str:
         if patched_image == self.image():
             patched_image = WORKFLOW_STATUS_PATCH_NOT_AVAILABLE
 
-        result = f"image: {self.repository}:{self.tag}\n" \
-                 f"\tscan status: {scan_status}\n" \
-                 f"\tscan date: {scan_date}\n" \
-                 f"\tscan task ID: {scan_task_id}\n" \
-                 f"\tpatch status: {patch_status}\n"
+        result = {
+            "image": f"{self.repository}:{self.tag}",
+            "scan_status": scan_status,
+            "scan_date": scan_date,
+            "scan_task_ID": scan_task_id,
+            "patch_status": patch_status
+        }
 
         if skipped_patch_reason != "":
-            result += f"\tskipped patch reason: {skipped_patch_reason}\n"
+            result["skipped_patch_reason"] = skipped_patch_reason
+
+        result["patch_date"] = patch_date
+        result["patch_task_ID"] = patch_task_id
+        result["last_patched_image"] = patched_image
+        result["workflow_type"] = workflow_type
 
-        result += f"\tpatch date: {patch_date}\n" \
-                  f"\tpatch task ID: {patch_task_id}\n" \
-                  f"\tlast patched image: {patched_image}\n" \
-                  f"\tworkflow type: {workflow_type}"
+        return result
+
+    def __str__(self) -> str:
+        status = self.get_status()
+        result = f"image: {status.repository}:{status.tag}\n" \
+                 f"\tscan status: {status.scan_status}\n" \
+                 f"\tscan date: {status.scan_date}\n" \
+                 f"\tscan task ID: {status.scan_task_id}\n" \
+                 f"\tpatch status: {status.patch_status}\n"
+
+        if hasattr(status, "skipped_patch_reason"):
+            result += f"\tskipped patch reason: {status.skipped_patch_reason}\n"
+
+        result += f"\tpatch date: {status.patch_date}\n" \
+                  f"\tpatch task ID: {status.patch_task_id}\n" \
+                  f"\tlast patched image: {status.patched_image}\n" \
+                  f"\tworkflow type: {status.workflow_type}"
 
         return result