diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 4a49ab7..569220d 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -17,9 +17,9 @@ jobs:
       - name: Run ECR Scan and Get SARIF Report
         uses: sartography/github-actions-library/wait-for-ecr-scan-and-get-sarif@main
         with:
-          repository_name: "my-ecr-repository"
-          image_tag: "latest"
-          aws_region: "us-west-2"
+          repository_name: "infr/testcloud2202"
+          image_tag: "main"
+          aws_region: "us-east-2"
           output_file: "report.sarif"
 
       - name: Upload SARIF report as artifact
diff --git a/wait-for-ecr-scan-and-get-sarif/pylib/wait_for_ecr_scan.py b/wait-for-ecr-scan-and-get-sarif/pylib/wait_for_ecr_scan.py
index 3eb61a5..2c89ff2 100644
--- a/wait-for-ecr-scan-and-get-sarif/pylib/wait_for_ecr_scan.py
+++ b/wait-for-ecr-scan-and-get-sarif/pylib/wait_for_ecr_scan.py
@@ -2,67 +2,50 @@ import json
 import boto3
 import time
 import sys
-
-
-# def wait_for_image_scan(repository_name, image_tag, region):
-#     client = boto3.client("ecr", region_name=region)
-#
-#     while True:
-#         response = client.describe_images(
-#             repositoryName=repository_name, imageIds=[{"imageTag": image_tag}]
-#         )
-#
-#         print(f"➡️ ➡️ ➡️  response: {response}")
-#         status = response["imageDetails"][0]["imageScanStatus"]["status"]
-#         print(f"Scan status: {status}")
-#
-#         if status == "COMPLETE":
-#             break
-#         elif status == "FAILED":
-#             raise Exception("Scan failed to complete")
-#         else:
-#             print("Still scanning, waiting for 30 seconds...")
-#             time.sleep(30)
+from botocore.exceptions import ClientError
 
 
 def wait_for_image_scan(repository_name, image_tag, region):
     client = boto3.client("ecr", region_name=region)
     response = None
 
-    while True:
+    max_retries = 10
+    retries = 0
+
+    while retries < max_retries:
         # maybe just check if this raises or not
-        response = client.describe_image_scan_findings(
-            repositoryName=repository_name, imageId={"imageTag": image_tag}
-        )
+        try:
+            response = client.describe_image_scan_findings(
+                repositoryName=repository_name, imageId={"imageTag": image_tag}
+            )
+        except ClientError as e:
+            if e.response['Error']['Code'] == 'ScanNotFoundException':
+                print(f"Scan not found for tag: {image_tag}. Retrying...")
+                retries += 1
+                time.sleep(30)
+                continue
+            else:
+                raise
 
         if "imageScanFindings" in response:
-            print("HIHIHI")
-        findings = response.get("imageScanFindings", {}).get("findings", [])
-        findings += response.get("imageScanFindings", {}).get("enhancedFindings", [])
-        print(f"Found {len(findings)} issues.")
-
-        if len(findings) > 0:
+            print(f"Scan found for repository {repository_name} and tag {image_tag}")
             break
 
+        # findings = response.get("imageScanFindings", {}).get("findings", [])
+        # findings += response.get("imageScanFindings", {}).get("enhancedFindings", [])
+        # print(f"Found {len(findings)} issues.")
+        #
+        # if len(findings) > 0:
+        #     break
+
         print("Still scanning, waiting for 30 seconds...")
         time.sleep(30)
-        # scan_status = response.get("imageScanStatus", {}).get("status")
-        #
-        # if not scan_status:
-        #     print(f"No image scan status found for tag: {image_tag}")
-        #     sys.exit(1)
-        #
-        # print(f"Scan status: {scan_status}")
-        #
-        # if scan_status == "COMPLETE":
-        #     print("Image scan complete!")
-        #     break
-        # elif scan_status == "FAILED":
-        #     print("Image scan failed!")
-        #     sys.exit(1)
-        # else:
+        retries += 1
 
-    return response
+    if retries == max_retries:
+        raise Exception("Max retries reached. Scan not found or incomplete.")
+    else:
+        return response
 
 
 if __name__ == "__main__":