From 80b611add092df7311416348a57a8dc8209af46e Mon Sep 17 00:00:00 2001
From: lingenj <jacob.van.lingen@moderne.io>
Date: Wed, 18 Dec 2024 16:16:34 +0100
Subject: [PATCH] work on it

---
 .../java/org/openrewrite/java/FindImage.java  | 38 ++++++++++-
 .../org/openrewrite/java/FindImageTest.java   | 65 ++++---------------
 2 files changed, 50 insertions(+), 53 deletions(-)

diff --git a/rewrite-java/src/main/java/org/openrewrite/java/FindImage.java b/rewrite-java/src/main/java/org/openrewrite/java/FindImage.java
index 5cd8c75aa78..a895ac84d35 100644
--- a/rewrite-java/src/main/java/org/openrewrite/java/FindImage.java
+++ b/rewrite-java/src/main/java/org/openrewrite/java/FindImage.java
@@ -15,14 +15,19 @@
  */
 package org.openrewrite.java;
 
+import lombok.EqualsAndHashCode;
+import lombok.Value;
 import org.jspecify.annotations.Nullable;
 import org.openrewrite.*;
 import org.openrewrite.java.table.ImageSourceFiles;
 import org.openrewrite.marker.SearchResult;
 import org.openrewrite.trait.Reference;
+import org.openrewrite.trait.SimpleTraitMatcher;
 
 import java.nio.file.Path;
 import java.util.Collection;
+import java.util.List;
+import java.util.Map;
 import java.util.stream.Collectors;
 
 // TODO: Remove this file, we will use the `FindDockerImageUses` in the rewrite-docker module
@@ -45,21 +50,52 @@ public TreeVisitor<?, ExecutionContext> getVisitor() {
 
             @Override
             public @Nullable Tree visit(@Nullable Tree tree, ExecutionContext ctx) {
-                // TODO improve: `if (sourceFile instanceof PlainText && references.size() > 1)` then all markers are set at beginning
+                Tree t = super.visit(tree, ctx);
 
                 if (tree instanceof SourceFileWithReferences) {
                     SourceFileWithReferences sourceFile = (SourceFileWithReferences) tree;
                     Path sourcePath = sourceFile.getSourcePath();
                     Collection<Reference> references = sourceFile.getReferences().findMatches(new ImageMatcher());
+
+                    /*new JavaIsoVisitor<ExecutionContext>() {
+                      vsit
+                    };
+
+                    new SimpleTraitMatcher<Reference>() {
+
+                        @Override
+                        protected @Nullable Reference test(Cursor cursor) {
+                            return null;
+                        }
+                    }.asVisitor().visit(sourceFile, 0);*/
+
+                    // TODO improve: `if (sourceFile instanceof PlainText && references.size() > 1)` then all markers are set at beginning
                     String value = references.stream()
                             .map(Reference::getValue)
                             .peek(it -> results.insertRow(ctx, new ImageSourceFiles.Row(sourcePath.toString(), tree.getClass().getSimpleName(), it)))
                             .sorted()
                             .collect(Collectors.joining("|"));
+                    /*System.out.println(tree);
+                    System.out.println(value);*/
                     return SearchResult.found(tree, value);
                 }
                 return tree;
             }
         };
     }
+
+    @Value
+    @EqualsAndHashCode(callSuper = false)
+    private static class ReferenceFindSearchResultVisitor extends TreeVisitor<Tree, ExecutionContext> {
+        Map<Tree, Reference> matches;
+
+        @Override
+        public Tree postVisit(Tree tree, ExecutionContext ctx) {
+            Reference reference = matches.get(tree);
+            if (reference != null && getCursor().equals(reference.getCursor())) {
+                return SearchResult.found(tree, reference.getValue());
+            }
+            return tree;
+        }
+    }
 }
diff --git a/rewrite-java/src/test/java/org/openrewrite/java/FindImageTest.java b/rewrite-java/src/test/java/org/openrewrite/java/FindImageTest.java
index d2fd481e9b6..81e468af507 100644
--- a/rewrite-java/src/test/java/org/openrewrite/java/FindImageTest.java
+++ b/rewrite-java/src/test/java/org/openrewrite/java/FindImageTest.java
@@ -17,9 +17,11 @@
 
 import org.junit.jupiter.api.Test;
 import org.openrewrite.DocumentExample;
+import org.openrewrite.java.table.ImageSourceFiles;
 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;
 import static org.openrewrite.yaml.Assertions.yaml;
 
@@ -35,64 +37,17 @@ public void defaults(RecipeSpec spec) {
     @Test
     void gitlabCIFile() {
         rewriteRun(
+          spec -> spec.recipe(new FindImage())
+            .dataTable(ImageSourceFiles.Row.class, rows -> {
+                assertThat(rows).hasSize(1);
+                assertThat(rows.get(0).getValue()).isEqualTo("maven:latest");
+            }),
           yaml(
             """
               image: maven:latest
-              
-              variables:
-                MAVEN_CLI_OPTS: "-s .m2/settings.xml --batch-mode"
-                MAVEN_OPTS: "-Dmaven.repo.local=.m2/repository"
-              
-              cache:
-                paths:
-                  - .m2/repository/
-                  - target/
-              
-              build:
-                stage: build
-                script:
-                  - mvn $MAVEN_CLI_OPTS compile
-              
-              test:
-                stage: test
-                script:
-                  - mvn $MAVEN_CLI_OPTS test
-              
-              deploy:
-                stage: deploy
-                script:
-                  - mvn $MAVEN_CLI_OPTS deploy
-                only:
-                  - master
               """,
             """
               ~~(maven:latest)~~>image: maven:latest
-              
-              variables:
-                MAVEN_CLI_OPTS: "-s .m2/settings.xml --batch-mode"
-                MAVEN_OPTS: "-Dmaven.repo.local=.m2/repository"
-              
-              cache:
-                paths:
-                  - .m2/repository/
-                  - target/
-              
-              build:
-                stage: build
-                script:
-                  - mvn $MAVEN_CLI_OPTS compile
-              
-              test:
-                stage: test
-                script:
-                  - mvn $MAVEN_CLI_OPTS test
-              
-              deploy:
-                stage: deploy
-                script:
-                  - mvn $MAVEN_CLI_OPTS deploy
-                only:
-                  - master
               """,
             spec -> spec.path(".gitlab-ci")
           )
@@ -102,6 +57,12 @@ void gitlabCIFile() {
     @Test
     void dockerFile() {
         rewriteRun(
+          spec -> spec.recipe(new FindImage())
+            .dataTable(ImageSourceFiles.Row.class, rows -> {
+                assertThat(rows).hasSize(2);
+                assertThat(rows.get(0).getValue()).isEqualTo("alpine:latest");
+                assertThat(rows.get(1).getValue()).isEqualTo("golang:1.7.3");
+            }),
           text(
             //language=Dockerfile
             """