Support replacing a custom method pattern too

openrewrite · Jan 24, 2024 · ccefdce · ccefdce
1 parent 5f838fb
commit ccefdce
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 6 deletions.
diff --git a/src/main/java/org/openrewrite/launchdarkly/RemoveBoolVariation.java b/src/main/java/org/openrewrite/launchdarkly/RemoveBoolVariation.java
@@ -18,7 +18,8 @@
 import lombok.EqualsAndHashCode;
 import lombok.Value;
 import org.openrewrite.*;
-import org.openrewrite.internal.lang.NonNull;
+import org.openrewrite.internal.StringUtils;
+import org.openrewrite.internal.lang.Nullable;
 import org.openrewrite.java.JavaVisitor;
 import org.openrewrite.java.MethodMatcher;
 import org.openrewrite.java.search.UsesMethod;
@@ -30,9 +31,14 @@
 import org.openrewrite.staticanalysis.RemoveUnusedPrivateFields;
 import org.openrewrite.staticanalysis.SimplifyConstantIfBranchExecution;
 
+import java.util.Optional;
+
 @Value
 @EqualsAndHashCode(callSuper = false)
 public class RemoveBoolVariation extends Recipe {
+
+    private static final String METHOD_PATTERN_BOOLVARIATION = "com.launchdarkly.sdk.server.LDClient boolVariation(String, com.launchdarkly.sdk.*, boolean)";
+
     @Override
     public String getDisplayName() {
         return "Remove `boolVariation` for feature key";
@@ -46,19 +52,25 @@ public String getDescription() {
     @Option(displayName = "Feature flag key",
             description = "The key of the feature flag to remove.",
             example = "flag-key-123abc")
-    @NonNull
     String featureKey;
 
     @Option(displayName = "Replacement value",
             description = "The value to replace the feature flag check with.",
             example = "true")
-    @NonNull
     Boolean replacementValue;
 
-    private static final MethodMatcher methodMatcher = new MethodMatcher("com.launchdarkly.sdk.server.LDClient boolVariation(String, com.launchdarkly.sdk.*, boolean)", true);
+    @Option(displayName = "Method pattern",
+            description = "A method pattern to match against. If not specified, will match `LDClient` `boolVariation`. " +
+                          "The first argument must be the feature key as `String`.",
+            example = METHOD_PATTERN_BOOLVARIATION,
+            required = false)
+    @Nullable
+    String methodPattern;
 
     @Override
     public TreeVisitor<?, ExecutionContext> getVisitor() {
+        String pattern = Optional.ofNullable(methodPattern).filter(StringUtils::isNotEmpty).orElse(METHOD_PATTERN_BOOLVARIATION);
+        final MethodMatcher methodMatcher = new MethodMatcher(pattern, true);
         JavaVisitor<ExecutionContext> visitor = new JavaVisitor<ExecutionContext>() {
             @Override
             public J visitMethodInvocation(J.MethodInvocation method, ExecutionContext ctx) {

diff --git a/src/test/java/org/openrewrite/launchdarkly/RemoveBoolVariationTest.java b/src/test/java/org/openrewrite/launchdarkly/RemoveBoolVariationTest.java
@@ -28,7 +28,7 @@ class RemoveBoolVariationTest implements RewriteTest {
 
     @Override
     public void defaults(RecipeSpec spec) {
-        spec.recipe(new RemoveBoolVariation("flag-key-123abc", true))
+        spec.recipe(new RemoveBoolVariation("flag-key-123abc", true, null))
           .parser(JavaParser.fromJavaVersion()
             .classpathFromResources(new InMemoryExecutionContext(), "launchdarkly-java-server-sdk-6"));
     }
@@ -107,7 +107,7 @@ void bar() {
     @Test
     void disablePermanently() {
         rewriteRun(
-          spec -> spec.recipe(new RemoveBoolVariation("flag-key-123abc", false)),
+          spec -> spec.recipe(new RemoveBoolVariation("flag-key-123abc", false, null)),
           // language=java
           java(
             """
@@ -250,4 +250,55 @@ void bar() {
           )
         );
     }
+
+    @Test
+    void customMethodPatternForWrapper() {
+        rewriteRun(
+          spec -> spec.recipe(new RemoveBoolVariation("flag-key-123abc", true, "com.acme.bank.CustomLaunchDarklyWrapper featureFlagEnabled(String, boolean)")),
+          // language=java
+          java(
+            """
+              package com.acme.bank;
+              
+              import com.launchdarkly.sdk.LDContext;
+              import com.launchdarkly.sdk.server.LDClient;
+              
+              public class CustomLaunchDarklyWrapper {
+                  private LDClient client = new LDClient("sdk-key-123abc");
+                  public boolean featureFlagEnabled(String key, boolean fallback) {
+                      LDContext context = null;
+                      return client.boolVariation(key, context, false);
+                  }
+              }
+              """
+          ),
+          // language=java
+          java(
+            """
+              import com.acme.bank.CustomLaunchDarklyWrapper;
+              class Foo {
+                  private CustomLaunchDarklyWrapper wrapper = new CustomLaunchDarklyWrapper();
+                  void bar() {
+                      if (wrapper.featureFlagEnabled("flag-key-123abc", false)) {
+                          // Application code to show the feature
+                          System.out.println("Feature is on");
+                      }
+                      else {
+                        // The code to run if the feature is off
+                          System.out.println("Feature is off");
+                      }
+                  }
+              }
+              """,
+            """
+              class Foo {
+                  void bar() {
+                      // Application code to show the feature
+                      System.out.println("Feature is on");
+                  }
+              }
+              """
+          )
+        );
+    }
 }