From b69bd42aad35804036dd8160722ee19d02eac1bc Mon Sep 17 00:00:00 2001
From: Tim te Beek <tim@moderne.io>
Date: Mon, 15 Jul 2024 11:31:58 +0200
Subject: [PATCH] Strip multi-stage names from tags (#2)

* Strip multi-stage names from tags

https://docs.docker.com/build/building/multi-stage/

* Drop explicit dependency
---
 .../openrewrite/docker/trait/Dockerfile.java  |  2 +-
 .../docker/FindDockerImagesUsedTest.java      | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/src/main/java/org/openrewrite/docker/trait/Dockerfile.java b/src/main/java/org/openrewrite/docker/trait/Dockerfile.java
index a116441..712542d 100644
--- a/src/main/java/org/openrewrite/docker/trait/Dockerfile.java
+++ b/src/main/java/org/openrewrite/docker/trait/Dockerfile.java
@@ -37,7 +37,7 @@ public List<DockerImageVersion> getFroms() {
                 String[] imageVersionStr = line.substring("FROM".length()).trim().split(":");
                 froms.add(new DockerImageVersion(
                         imageVersionStr[0],
-                        imageVersionStr.length > 1 ? imageVersionStr[1] : null
+                        imageVersionStr.length > 1 ? imageVersionStr[1].split(" ")[0] : null
                 ));
             }
         }
diff --git a/src/test/java/org/openrewrite/docker/FindDockerImagesUsedTest.java b/src/test/java/org/openrewrite/docker/FindDockerImagesUsedTest.java
index e26c0c0..47cfbb2 100644
--- a/src/test/java/org/openrewrite/docker/FindDockerImagesUsedTest.java
+++ b/src/test/java/org/openrewrite/docker/FindDockerImagesUsedTest.java
@@ -18,9 +18,11 @@
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
 import org.openrewrite.docker.search.FindDockerImageUses;
+import org.openrewrite.docker.table.DockerBaseImages;
 import org.openrewrite.test.RecipeSpec;
 import org.openrewrite.test.RewriteTest;
 
+import static org.assertj.core.api.Assertions.assertThat;
 import static org.openrewrite.test.SourceSpecs.text;
 
 class FindDockerImagesUsedTest implements RewriteTest {
@@ -52,4 +54,28 @@ void dockerfile() {
           )
         );
     }
+
+    @Test
+    void multistageDockerfile() {
+        rewriteRun(
+          spec -> spec.dataTable(DockerBaseImages.Row.class, rows -> assertThat(rows)
+            .containsOnly(new DockerBaseImages.Row("nvidia/cuda", "11.8.0-cudnn8-devel-ubuntu20.04"))),
+          text(
+            //language=Dockerfile
+            """
+              FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04 AS base
+              LABEL maintainer="Hugging Face"
+              ARG DEBIAN_FRONTEND=noninteractive
+              SHELL ["sh", "-lc"]
+              """,
+            """
+              ~~(nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04)~~>FROM nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04 AS base
+              LABEL maintainer="Hugging Face"
+              ARG DEBIAN_FRONTEND=noninteractive
+              SHELL ["sh", "-lc"]
+              """,
+            spec -> spec.path("Dockerfile")
+          )
+        );
+    }
 }