From be854ee003305871e4a9505c0f6ee3a9820d0e08 Mon Sep 17 00:00:00 2001
From: Avi Biton <abiton@redhat.com>
Date: Sun, 17 Dec 2023 15:41:47 +0200
Subject: [PATCH] chore(RHTAPWATCH-682): add the ability to extract images from
 snapshot

add the ability to extract images from snapshot
add tests

Signed-off-by: Avi Biton <abiton@redhat.com>
---
 tests/test_rpm_verifier.py  | 53 ++++++++++++++++++++++++++++++++++++-
 verify_rpms/rpm_verifier.py | 13 ++++++---
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/tests/test_rpm_verifier.py b/tests/test_rpm_verifier.py
index 55bd0d1..5f91bea 100644
--- a/tests/test_rpm_verifier.py
+++ b/tests/test_rpm_verifier.py
@@ -1,5 +1,4 @@
 """Test rpm_verifier.py end-to-end"""
-
 from pathlib import Path
 from textwrap import dedent
 from unittest.mock import MagicMock, call, create_autospec, sentinel
@@ -14,6 +13,7 @@
     generate_output,
     get_rpmdb,
     get_unsigned_rpms,
+    parse_image_input,
 )
 
 
@@ -130,6 +130,57 @@ def test_get_unsigned_rpms(test_input: list[str], expected: list[str]) -> None:
     assert result == expected
 
 
+@pytest.mark.parametrize(
+    "test_input,expected",
+    [
+        pytest.param(
+            dedent(
+                """
+                {
+                    "spec": {
+                        "application": "test",
+                        "components": [
+                            {
+                                "containerImage": "quay.io/container-image@sha256:123"
+                            },
+                            {
+                                "containerImage": "quay.io/container-image@sha256:456"
+                            },      
+                            {
+                                "containerImage": "quay.io/container-image@sha256:789"
+                            }
+                        ]
+                    }
+                }
+                """
+            ).strip(),
+            [
+                "quay.io/container-image@sha256:123",
+                "quay.io/container-image@sha256:456",
+                "quay.io/container-image@sha256:789",
+            ],
+            id="snapshot test",
+        ),
+        pytest.param(
+            "quay.io/container-image@sha256:123",
+            ["quay.io/container-image@sha256:123"],
+            id="single container",
+        ),
+    ],
+)
+def test_parse_image_input(test_input: str, expected: list[str]) -> None:
+    """Test parse_image_input"""
+    result = parse_image_input(image_input=test_input)
+    assert result == expected
+
+
+def test_parse_image_input_exception() -> None:
+    """Test parse_image_input throws exception"""
+    test_input = '{"apiVersion": "appstudio.redhat.com/v1alpha1"}'
+    with pytest.raises(KeyError):
+        parse_image_input(image_input=test_input)
+
+
 class TestImageProcessor:
     """Test ImageProcessor's callable"""
 
diff --git a/verify_rpms/rpm_verifier.py b/verify_rpms/rpm_verifier.py
index 42249b0..eb0d08d 100644
--- a/verify_rpms/rpm_verifier.py
+++ b/verify_rpms/rpm_verifier.py
@@ -6,6 +6,7 @@
 import tempfile
 from concurrent.futures import ThreadPoolExecutor
 from dataclasses import dataclass
+from json import JSONDecodeError, loads
 from pathlib import Path
 from subprocess import run
 from typing import Callable, Iterable
@@ -97,9 +98,15 @@ def parse_image_input(image_input: str) -> list[str]:
     Try parsing as json and extract the images.
     If failing to parse as json, assume it's an image reference as one-element list.
     """
-    # Currently assuming a single image. To be replaced by deciding whether this is an
-    # image or a snapshot json, and parsing accordingly
-    return [image_input]
+    try:
+        snapshot = loads(s=image_input)
+    except JSONDecodeError:
+        return [image_input]
+    components = snapshot["spec"]["components"]
+    container_images: list[str] = []
+    for component in components:
+        container_images.append(component["containerImage"])
+    return container_images
 
 
 @dataclass(frozen=True)