From 3575f7a268221b13637ffc0345328334776f3593 Mon Sep 17 00:00:00 2001
From: guoje <guoje@amazon.com>
Date: Fri, 11 Mar 2022 14:04:28 -0800
Subject: [PATCH] fix: remove test docker image from device

---
 .../testing/features/DockerSteps.java         | 22 +++++++++++++++----
 .../greengrass/features/docker.feature        |  2 +-
 2 files changed, 19 insertions(+), 5 deletions(-)

diff --git a/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/java/com/aws/greengrass/testing/features/DockerSteps.java b/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/java/com/aws/greengrass/testing/features/DockerSteps.java
index 6fe60173..f0cf8514 100644
--- a/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/java/com/aws/greengrass/testing/features/DockerSteps.java
+++ b/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/java/com/aws/greengrass/testing/features/DockerSteps.java
@@ -39,11 +39,13 @@ public class DockerSteps {
      *
      * @param image fully qualified image name in <code>name:tag</code> representation
      */
-    @Given("the docker image {word} does not exist on the device")
+    @Given("deleted the docker image {word} if exists on the device")
     public void checkDockerImageIsMissing(String image) {
         Predicate<Set<String>> predicate = Set::isEmpty;
-        checkDockerImagePresence(image, predicate.negate(),
-                "The image " + image + " is already on the device. Please remove the image and try again.");
+        if (isDockerImagePresence(image, predicate.negate())) {
+            removeDockerImage(image);
+        }
+
     }
 
     @Then("I can check that the docker image {word} exists on the device")
@@ -65,18 +67,30 @@ public void removeDockerImage(String image) {
         LOGGER.debug("Removed docker image {}: {}", image, result);
     }
 
-    private void checkDockerImagePresence(String image, Predicate<Set<String>> validity, String message) {
+    private boolean isDockerImagePresence(String image, Predicate<Set<String>> validity) {
         // This could be improved by using the API on a local host.
+
         Set<String> parts = new HashSet<>(Arrays.stream(image.split(":")).collect(Collectors.toSet()));
         String[] result = platform.commands().executeToString(CommandInput.builder()
                 .line("docker").addArgs("images")
                 .build())
                 .split("\\r?\\n");
+
         Arrays.stream(result)
                 .map(String::trim)
                 .flatMap(line -> Arrays.stream(line.split("\\s+")))
                 .forEach(parts::remove);
+
         if (!validity.test(parts)) {
+            return false;
+        }
+
+        return true;
+    }
+
+    private void checkDockerImagePresence(String image, Predicate<Set<String>> validity, String message) {
+
+        if (!isDockerImagePresence(image, validity)) {
             throw new IllegalStateException(message);
         }
     }
diff --git a/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/resources/greengrass/features/docker.feature b/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/resources/greengrass/features/docker.feature
index be0027d4..94e9ff77 100644
--- a/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/resources/greengrass/features/docker.feature
+++ b/aws-greengrass-testing-features/aws-greengrass-testing-features-docker/src/main/resources/greengrass/features/docker.feature
@@ -6,7 +6,7 @@ Feature: Greengrass V2 Docker Component
 
   @Docker @IDT
   Scenario: I can deploy Docker containers as Greengrass Components
-    Given the docker image amazon/amazon-ec2-metadata-mock:v1.9.0 does not exist on the device
+    Given deleted the docker image amazon/amazon-ec2-metadata-mock:v1.9.0 if exists on the device
     And I create a Greengrass deployment with components
       | DockerHubAmazonContainer | classpath:/greengrass/component/recipes/DockerHubAmazonContainer.yaml |
     When I deploy the Greengrass deployment configuration