diff --git a/java/README.md b/java/README.md
index 5931fd9aa..395bc3139 100644
--- a/java/README.md
+++ b/java/README.md
@@ -65,6 +65,13 @@ The properties file has the following template:
       },
       ...
     ],
+  "testMethodProperties": [
+    {
+      "methodName": "when",
+      "argumentIndex": 0
+    },
+    ...
+  ],
   "enumProperties":
     [
       {
@@ -106,7 +113,13 @@ public void some_unit_test() { ... }
 
 when `IsTreated` is `true`, and will be deleted completely when `IsTreated` is `false`.
 
-An optional top-level field is `enumProperties`.
+An optional top-level field is `testMethodProperties`.
+
+Within that, there is an array of JSON objects, having the required fields `methodName` and `argumentIndex`. They both behave the same as the fields with the same name in `methodProperties`.
+
+What this field does, is that if we find one of the `methodProperties` fields inside a method that matches one of the methods in `testMethodProperties`, we remove that method. This is useful for removing `mock()` wrappers or `assert()` calls that are no longer useful after a flag is cleaned up.
+
+Another optional top-level field is `enumProperties`.
 Within that, there is an array of JSON objects, having the required fields `enumName` and `argumentIndex`.
 
 What this field does, is if you specify an enum class name, Piranha will remove enum constants that have a constructor with a string argument that matches your `FlagName` value, along with their usages.
diff --git a/java/piranha/config/properties.json b/java/piranha/config/properties.json
index 983765365..95332b775 100644
--- a/java/piranha/config/properties.json
+++ b/java/piranha/config/properties.json
@@ -38,6 +38,20 @@
       "argumentIndex": 0
     }
   ],
+  "testMethodProperties": [
+    {
+      "methodName": "mock",
+      "argumentIndex": 0
+    },
+    {
+      "methodName": "accept",
+      "argumentIndex": 0
+    },
+    {
+      "methodName": "expect",
+      "argumentIndex": 1
+    }
+  ],
   "linkURL": "<provide_your_url>",
   "annotations": [
     "ToggleTesting",
diff --git a/java/piranha/src/main/java/com/uber/piranha/XPFlagCleaner.java b/java/piranha/src/main/java/com/uber/piranha/XPFlagCleaner.java
index 771709ba8..4e882adce 100644
--- a/java/piranha/src/main/java/com/uber/piranha/XPFlagCleaner.java
+++ b/java/piranha/src/main/java/com/uber/piranha/XPFlagCleaner.java
@@ -58,7 +58,8 @@
 import com.uber.piranha.config.Config;
 import com.uber.piranha.config.PiranhaConfigurationException;
 import com.uber.piranha.config.PiranhaEnumRecord;
-import com.uber.piranha.config.PiranhaMethodRecord;
+import com.uber.piranha.config.PiranhaFlagMethodRecord;
+import com.uber.piranha.config.PiranhaTestMethodRecord;
 import com.uber.piranha.testannotations.AnnotationArgument;
 import com.uber.piranha.testannotations.ResolvedTestAnnotation;
 import java.util.ArrayList;
@@ -344,7 +345,7 @@ private API getXPAPI(ExpressionTree et, VisitorState state) {
       }
       MemberSelectTree mst = (MemberSelectTree) mit.getMethodSelect();
       String methodName = mst.getIdentifier().toString();
-      ImmutableCollection<PiranhaMethodRecord> methodRecords =
+      ImmutableCollection<PiranhaFlagMethodRecord> methodRecords =
           this.config.getMethodRecordsForName(methodName);
       if (methodRecords.size() > 0) {
         return getXPAPI(mit, state, methodRecords);
@@ -356,8 +357,8 @@ private API getXPAPI(ExpressionTree et, VisitorState state) {
   private API getXPAPI(
       MethodInvocationTree mit,
       VisitorState state,
-      ImmutableCollection<PiranhaMethodRecord> methodRecordsForName) {
-    for (PiranhaMethodRecord methodRecord : methodRecordsForName) {
+      ImmutableCollection<PiranhaFlagMethodRecord> methodRecordsForName) {
+    for (PiranhaFlagMethodRecord methodRecord : methodRecordsForName) {
       // when argumentIndex is specified, if mit's argument at argIndex doesn't match xpFlagName,
       // skip to next method property map
       Optional<Integer> optionalArgumentIdx = methodRecord.getArgumentIdx();
@@ -880,6 +881,31 @@ public Description matchBinary(BinaryTree tree, VisitorState state) {
     return Description.NO_MATCH;
   }
 
+  private boolean matchTestMethod(MethodInvocationTree methodTree, VisitorState state) {
+    Symbol receiverSymbol = ASTHelpers.getSymbol(methodTree.getMethodSelect());
+    String methodName = receiverSymbol.getSimpleName().toString();
+    for (PiranhaTestMethodRecord methodRecord : config.getTestMethodRecordsForName(methodName)) {
+      Optional<Integer> argumentIdx = methodRecord.getArgumentIdx();
+      if (argumentIdx.isPresent()) {
+        if (methodTree.getArguments().size() > argumentIdx.get()) {
+          ExpressionTree argTree = methodTree.getArguments().get(argumentIdx.get());
+          API api = getXPAPI(argTree, state);
+          if (api != API.UNKNOWN) {
+            return true;
+          }
+        }
+      } else {
+        for (ExpressionTree argTree : methodTree.getArguments()) {
+          API api = getXPAPI(argTree, state);
+          if (api != API.UNKNOWN) {
+            return true;
+          }
+        }
+      }
+    }
+    return false;
+  }
+
   @Override
   public Description matchExpressionStatement(ExpressionStatementTree tree, VisitorState state) {
     if (shouldSkip(state)) return Description.NO_MATCH;
@@ -887,12 +913,29 @@ public Description matchExpressionStatement(ExpressionStatementTree tree, Visito
       return Description.NO_MATCH;
     }
 
+    boolean updateCode = false;
+
     if (tree.getExpression().getKind().equals(Kind.METHOD_INVOCATION)) {
       MethodInvocationTree mit = (MethodInvocationTree) tree.getExpression();
+      ExpressionTree receiver = ASTHelpers.getReceiver(mit);
+      if (receiver == null) {
+        if (matchTestMethod(mit, state)) {
+          updateCode = true;
+        }
+      } else if (receiver.getKind() == Kind.METHOD_INVOCATION) {
+        if (matchTestMethod((MethodInvocationTree) receiver, state)) {
+          updateCode = true;
+        }
+      }
+
       API api = getXPAPI(mit, state);
       if (api.equals(API.DELETE_METHOD)
           || api.equals(API.SET_TREATED)
           || api.equals(API.SET_CONTROL)) {
+        updateCode = true;
+      }
+
+      if (updateCode) {
         Description.Builder builder = buildDescription(tree);
         SuggestedFix.Builder fixBuilder = SuggestedFix.builder();
         fixBuilder.delete(tree);
@@ -1192,7 +1235,7 @@ private void recursiveScanTestMethodStats(
         // only when the flag name matches, and we want to verify that no calls are being made to
         // set
         // unrelated flags (i.e. count them in counters.allSetters).
-        for (PiranhaMethodRecord methodRecord : config.getMethodRecordsForName(methodName)) {
+        for (PiranhaFlagMethodRecord methodRecord : config.getMethodRecordsForName(methodName)) {
           if (methodRecord.getApiType().equals(XPFlagCleaner.API.SET_TREATED)) {
             counters.allSetters += 1;
             // If the test is asking for the flag in treated condition, but we are setting it to
diff --git a/java/piranha/src/main/java/com/uber/piranha/config/Config.java b/java/piranha/src/main/java/com/uber/piranha/config/Config.java
index 313b4163b..cdfc8a18b 100644
--- a/java/piranha/src/main/java/com/uber/piranha/config/Config.java
+++ b/java/piranha/src/main/java/com/uber/piranha/config/Config.java
@@ -41,7 +41,8 @@ public final class Config {
   /* Names of top-level fields within properties.json */
   private static final String LINK_URL_KEY = "linkURL";
   private static final String ANNOTATIONS_KEY = "annotations";
-  private static final String METHODS_KEY = "methodProperties";
+  private static final String FLAG_METHODS_KEY = "methodProperties";
+  private static final String TEST_METHODS_KEY = "testMethodProperties";
   private static final String CLEANUP_OPTS_KEY = "cleanupOptions";
   private static final String ENUMS_KEY = "enumProperties";
 
@@ -64,13 +65,22 @@ public final class Config {
   private static final boolean DEFAULT_TESTS_CLEAN_BY_SETTERS_IGNORE_OTHERS = false;
 
   /**
-   * configMethodsMap is a map where key is method name and value is a list where each item in the
-   * list is a map that corresponds to each method property from properties.json. In most cases, the
-   * list would have only one element. But if someone reuses the same method name with different
+   * configMethodProperties is a map where key is method name and value is a list where each item in
+   * the list is a map that corresponds to each method property from properties.json. In most cases,
+   * the list would have only one element. But if someone reuses the same method name with different
    * returnType/receiverType/argumentIndex, the list would have each method property map as one
    * element.
    */
-  private final ImmutableMultimap<String, PiranhaMethodRecord> configMethodProperties;
+  private final ImmutableMultimap<String, PiranhaFlagMethodRecord> configMethodProperties;
+
+  /**
+   * configTestMethodProperties is a map where key is method name and value is a list where each
+   * item in the list is a map that corresponds to each test method property from properties.json.
+   * In most cases, the list would have only one element. But if someone reuses the same method name
+   * with different returnType/receiverType/argumentIndex, the list would have each method property
+   * map as one element.
+   */
+  private final ImmutableMultimap<String, PiranhaTestMethodRecord> configTestMethodProperties;
 
   /**
    * configEnumProperties is a map where key is enum name and value is a list where each item in the
@@ -99,12 +109,14 @@ public final class Config {
   // Constructor is private, a Config object can be generated using the class' static methods,
   // in particular Config.fromJSONFile([properties.json])
   private Config(
-      ImmutableMultimap<String, PiranhaMethodRecord> configMethodProperties,
+      ImmutableMultimap<String, PiranhaFlagMethodRecord> configMethodProperties,
+      ImmutableMultimap<String, PiranhaTestMethodRecord> configTestMethodProperties,
       ImmutableMultimap<String, PiranhaEnumRecord> configEnumProperties,
       TestAnnotationResolver testAnnotationResolver,
       ImmutableMap<String, Object> cleanupOptions,
       String linkURL) {
     this.configMethodProperties = configMethodProperties;
+    this.configTestMethodProperties = configTestMethodProperties;
     this.configEnumProperties = configEnumProperties;
     this.testAnnotationResolver = testAnnotationResolver;
     this.cleanupOptions = cleanupOptions;
@@ -115,15 +127,29 @@ private Config(
    * Return all configuration method records matching a given method name.
    *
    * @param methodName the method name to search
-   * @return A collection of {@link PiranhaMethodRecord} objects, representing each method
+   * @return A collection of {@link PiranhaFlagMethodRecord} objects, representing each method
    *     definition in the piranha json configuration file matching {@code methodName}.
    */
-  public ImmutableCollection<PiranhaMethodRecord> getMethodRecordsForName(String methodName) {
+  public ImmutableCollection<PiranhaFlagMethodRecord> getMethodRecordsForName(String methodName) {
     return configMethodProperties.containsKey(methodName)
         ? configMethodProperties.get(methodName)
         : ImmutableSet.of();
   }
 
+  /**
+   * Return all configuration test method records matching a given method name.
+   *
+   * @param methodName the method name to search
+   * @return A collection of {@link PiranhaTestMethodRecord} objects, representing each method
+   *     definition in the piranha json configuration file matching {@code methodName}.
+   */
+  public ImmutableCollection<PiranhaTestMethodRecord> getTestMethodRecordsForName(
+      String methodName) {
+    return configTestMethodProperties.containsKey(methodName)
+        ? configTestMethodProperties.get(methodName)
+        : ImmutableSet.of();
+  }
+
   /**
    * Returns whether any configuration enum records exist. Useful for skipping logic if enum
    * properties are not configured.
@@ -266,7 +292,9 @@ public static Config fromJSONFile(String configFile, boolean isArgumentIndexOpti
       }
 
       String linkURL = DEFAULT_PIRANHA_URL;
-      ImmutableMultimap.Builder<String, PiranhaMethodRecord> methodsBuilder =
+      ImmutableMultimap.Builder<String, PiranhaFlagMethodRecord> methodsBuilder =
+          ImmutableMultimap.builder();
+      ImmutableMultimap.Builder<String, PiranhaTestMethodRecord> testMethodsBuilder =
           ImmutableMultimap.builder();
       ImmutableMultimap.Builder<String, PiranhaEnumRecord> enumsBuilder =
           ImmutableMultimap.builder();
@@ -294,17 +322,26 @@ public static Config fromJSONFile(String configFile, boolean isArgumentIndexOpti
           }
         }
       }
-      if (propertiesJson.get(METHODS_KEY) != null) {
+      if (propertiesJson.get(FLAG_METHODS_KEY) != null) {
         for (Map<String, Object> methodProperty :
-            (List<Map<String, Object>>) propertiesJson.get(METHODS_KEY)) {
-          PiranhaMethodRecord methodRecord =
-              PiranhaMethodRecord.parseFromJSONPropertyEntryMap(
+            (List<Map<String, Object>>) propertiesJson.get(FLAG_METHODS_KEY)) {
+          PiranhaFlagMethodRecord methodRecord =
+              PiranhaFlagMethodRecord.parseFromJSONPropertyEntryMap(
                   methodProperty, isArgumentIndexOptional);
           methodsBuilder.put(methodRecord.getMethodName(), methodRecord);
         }
       } else {
         throw new PiranhaConfigurationException("methodProperties not found, required.");
       }
+      if (propertiesJson.get(TEST_METHODS_KEY) != null) {
+        for (Map<String, Object> methodProperty :
+            (List<Map<String, Object>>) propertiesJson.get(TEST_METHODS_KEY)) {
+          PiranhaTestMethodRecord methodRecord =
+              PiranhaTestMethodRecord.parseFromJSONPropertyEntryMap(
+                  methodProperty, isArgumentIndexOptional);
+          testMethodsBuilder.put(methodRecord.getMethodName(), methodRecord);
+        }
+      }
       if (propertiesJson.get(ENUMS_KEY) != null) {
         for (Map<String, Object> enumProperty :
             (List<Map<String, Object>>) propertiesJson.get(ENUMS_KEY)) {
@@ -327,6 +364,7 @@ public static Config fromJSONFile(String configFile, boolean isArgumentIndexOpti
       }
       return new Config(
           methodsBuilder.build(),
+          testMethodsBuilder.build(),
           enumsBuilder.build(),
           annotationResolverBuilder.build(),
           cleanupOptionsBuilder.build(),
@@ -365,6 +403,7 @@ public static Config fromJSONFile(String configFile, boolean isArgumentIndexOpti
    */
   public static Config emptyConfig() {
     return new Config(
+        ImmutableMultimap.of(),
         ImmutableMultimap.of(),
         ImmutableMultimap.of(),
         TestAnnotationResolver.builder().build(),
diff --git a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaFlagMethodRecord.java b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaFlagMethodRecord.java
new file mode 100644
index 000000000..286f0b8d7
--- /dev/null
+++ b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaFlagMethodRecord.java
@@ -0,0 +1,120 @@
+/**
+ * Copyright (c) 2021 Uber Technologies, Inc.
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.uber.piranha.config;
+
+import static com.uber.piranha.config.PiranhaMethodRecord.parseArgumentIndexFromJSON;
+import static com.uber.piranha.config.PiranhaMethodRecord.parseFlagTypeFromJSON;
+import static com.uber.piranha.config.PiranhaMethodRecord.parseMethodNameFromJSON;
+import static com.uber.piranha.config.PiranhaRecord.getValueStringFromMap;
+
+import com.google.common.collect.ImmutableMap;
+import com.uber.piranha.XPFlagCleaner;
+import java.util.Map;
+import java.util.Optional;
+
+/** A class representing a flag method configuration record from properties.json */
+public final class PiranhaFlagMethodRecord implements PiranhaMethodRecord {
+  // Used for generating exception messages
+  private static final String PROPERTY_NAME = "methodProperty";
+
+  /**
+   * Holds the mapping of flagType string to API. Eg: "treated" -> API.IS_TREATED. Is initialized
+   * once and then accessed without updating.
+   */
+  private static final ImmutableMap<String, XPFlagCleaner.API> flagTypeToAPIMap =
+      initializeFlagTypeToAPIMap();
+
+  private final String methodName;
+  private final XPFlagCleaner.API apiType;
+  private final Optional<Integer> argumentIdx;
+  private final Optional<String> receiverType;
+  private final Optional<String> returnType;
+
+  PiranhaFlagMethodRecord(
+      String methodName,
+      String flagTypeString,
+      Optional<Integer> argumentIdx,
+      Optional<String> receiverType,
+      Optional<String> returnType) {
+    this.methodName = methodName;
+    this.apiType = flagTypeToAPIMap.getOrDefault(flagTypeString, XPFlagCleaner.API.UNKNOWN);
+    this.argumentIdx = argumentIdx;
+    this.receiverType = receiverType;
+    this.returnType = returnType;
+  }
+
+  @Override
+  public String getMethodName() {
+    return methodName;
+  }
+
+  public XPFlagCleaner.API getApiType() {
+    return apiType;
+  }
+
+  @Override
+  public Optional<Integer> getArgumentIdx() {
+    return argumentIdx;
+  }
+
+  @Override
+  public Optional<String> getReceiverType() {
+    return receiverType;
+  }
+
+  @Override
+  public Optional<String> getReturnType() {
+    return returnType;
+  }
+
+  private static ImmutableMap<String, XPFlagCleaner.API> initializeFlagTypeToAPIMap() {
+    ImmutableMap.Builder<String, XPFlagCleaner.API> builder = new ImmutableMap.Builder<>();
+    builder.put("treated", XPFlagCleaner.API.IS_TREATED);
+    builder.put("control", XPFlagCleaner.API.IS_CONTROL);
+    builder.put("set_treated", XPFlagCleaner.API.SET_TREATED);
+    builder.put("set_control", XPFlagCleaner.API.SET_CONTROL);
+    builder.put("empty", XPFlagCleaner.API.DELETE_METHOD);
+    builder.put("treatmentGroup", XPFlagCleaner.API.IS_TREATMENT_GROUP_CHECK);
+    return builder.build();
+  }
+
+  /**
+   * Parse the entry for a single method from piranha.json that has been previously decoded into a
+   * map
+   *
+   * @param methodPropertyEntry The decoded json entry (as a Map of property names to values)
+   * @param isArgumentIndexOptional Whether argumentIdx should be treated as optional
+   * @return A PiranhaFlagMethodRecord corresponding to the given map/json record.
+   * @throws PiranhaConfigurationException if there was any issue reading or parsing the
+   *     configuration file.
+   */
+  static PiranhaFlagMethodRecord parseFromJSONPropertyEntryMap(
+      Map<String, Object> methodPropertyEntry, boolean isArgumentIndexOptional)
+      throws PiranhaConfigurationException {
+
+    String methodName = parseMethodNameFromJSON(PROPERTY_NAME, methodPropertyEntry);
+    String flagType = parseFlagTypeFromJSON(PROPERTY_NAME, methodPropertyEntry);
+    Integer argumentIndexInteger =
+        parseArgumentIndexFromJSON(PROPERTY_NAME, methodPropertyEntry, isArgumentIndexOptional);
+    String receiverType = getValueStringFromMap(methodPropertyEntry, RECEIVER_TYPE_STRING);
+    String returnType = getValueStringFromMap(methodPropertyEntry, RETURN_TYPE_STRING);
+
+    return new PiranhaFlagMethodRecord(
+        methodName,
+        flagType,
+        Optional.ofNullable(argumentIndexInteger),
+        Optional.ofNullable(receiverType),
+        Optional.ofNullable(returnType));
+  }
+}
diff --git a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaMethodRecord.java b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaMethodRecord.java
index adae61828..a19e8a300 100644
--- a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaMethodRecord.java
+++ b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaMethodRecord.java
@@ -16,122 +16,70 @@
 import static com.uber.piranha.config.PiranhaRecord.getArgumentIndexFromMap;
 import static com.uber.piranha.config.PiranhaRecord.getValueStringFromMap;
 
-import com.google.common.collect.ImmutableMap;
-import com.uber.piranha.XPFlagCleaner;
 import java.util.Map;
 import java.util.Optional;
+import javax.annotation.Nullable;
 
-/** A class representing a method configuration record from properties.json */
-public final class PiranhaMethodRecord {
-
+/** An interface representing a method configuration record from properties.json */
+public interface PiranhaMethodRecord {
   // Allowed fields for a method property in the config file.
-  // Entered under the top-level "methodProperties" in properties.json.
+  // Entered under the top-level "methodProperties" or "testMethodProperties" in properties.json.
   // By default, the flagType, methodName and argumentIndex fields are mandatory.
   // The returnType and receiverType fields are optional.
-  private static final String FLAG_TYPE_KEY = "flagType";
-  private static final String METHOD_NAME_KEY = "methodName";
-  private static final String ARGUMENT_INDEX_KEY = "argumentIndex";
-  private static final String RETURN_TYPE_STRING = "returnType";
-  private static final String RECEIVER_TYPE_STRING = "receiverType";
+  String FLAG_TYPE_KEY = "flagType";
+  String METHOD_NAME_KEY = "methodName";
+  String ARGUMENT_INDEX_KEY = "argumentIndex";
+  String RETURN_TYPE_STRING = "returnType";
+  String RECEIVER_TYPE_STRING = "receiverType";
 
-  /**
-   * Holds the mapping of flagType string to API. Eg: "treated" -> API.IS_TREATED. Is initialized
-   * once and then accessed without updating.
-   */
-  private static final ImmutableMap<String, XPFlagCleaner.API> flagTypeToAPIMap =
-      initializeFlagTypeToAPIMap();
+  String getMethodName();
 
-  private final String methodName;
-  private final XPFlagCleaner.API apiType;
-  private final Optional<Integer> argumentIdx;
-  private final Optional<String> receiverType;
-  private final Optional<String> returnType;
+  Optional<Integer> getArgumentIdx();
 
-  PiranhaMethodRecord(
-      String methodName,
-      String flagTypeString,
-      Optional<Integer> argumentIdx,
-      Optional<String> receiverType,
-      Optional<String> returnType) {
-    this.methodName = methodName;
-    this.apiType = flagTypeToAPIMap.getOrDefault(flagTypeString, XPFlagCleaner.API.UNKNOWN);
-    this.argumentIdx = argumentIdx;
-    this.receiverType = receiverType;
-    this.returnType = returnType;
-  }
+  Optional<String> getReceiverType();
 
-  public String getMethodName() {
-    return methodName;
-  }
+  Optional<String> getReturnType();
 
-  public XPFlagCleaner.API getApiType() {
-    return apiType;
-  }
-
-  public Optional<Integer> getArgumentIdx() {
-    return argumentIdx;
-  }
-
-  public Optional<String> getReceiverType() {
-    return receiverType;
-  }
-
-  public Optional<String> getReturnType() {
-    return returnType;
+  static String parseMethodNameFromJSON(
+      String propertyName, Map<String, Object> methodPropertyEntry)
+      throws PiranhaConfigurationException {
+    String methodName = getValueStringFromMap(methodPropertyEntry, METHOD_NAME_KEY);
+    if (methodName == null) {
+      throw new PiranhaConfigurationException(
+          propertyName + " is missing mandatory methodName field. Check:\n" + methodPropertyEntry);
+    }
+    return methodName;
   }
 
-  private static ImmutableMap<String, XPFlagCleaner.API> initializeFlagTypeToAPIMap() {
-    ImmutableMap.Builder<String, XPFlagCleaner.API> builder = new ImmutableMap.Builder<>();
-    builder.put("treated", XPFlagCleaner.API.IS_TREATED);
-    builder.put("control", XPFlagCleaner.API.IS_CONTROL);
-    builder.put("set_treated", XPFlagCleaner.API.SET_TREATED);
-    builder.put("set_control", XPFlagCleaner.API.SET_CONTROL);
-    builder.put("empty", XPFlagCleaner.API.DELETE_METHOD);
-    builder.put("treatmentGroup", XPFlagCleaner.API.IS_TREATMENT_GROUP_CHECK);
-    return builder.build();
+  static String parseFlagTypeFromJSON(String propertyName, Map<String, Object> methodPropertyEntry)
+      throws PiranhaConfigurationException {
+    String flagType = getValueStringFromMap(methodPropertyEntry, FLAG_TYPE_KEY);
+    if (flagType == null) {
+      throw new PiranhaConfigurationException(
+          propertyName + " is missing mandatory flagType field. Check:\n" + methodPropertyEntry);
+    }
+    return flagType;
   }
 
-  /**
-   * Parse the entry for a single method from piranha.json that has been previously decoded into a
-   * map
-   *
-   * @param methodPropertyEntry The decoded json entry (as a Map of property names to values)
-   * @param isArgumentIndexOptional Whether argumentIdx should be treated as optional
-   * @return A PiranhaMethodRecord corresponding to the given map/json record.
-   * @throws PiranhaConfigurationException if there was any issue reading or parsing the
-   *     configuration file.
-   */
-  static PiranhaMethodRecord parseFromJSONPropertyEntryMap(
-      Map<String, Object> methodPropertyEntry, boolean isArgumentIndexOptional)
+  @Nullable
+  static Integer parseArgumentIndexFromJSON(
+      String propertyName, Map<String, Object> methodPropertyEntry, boolean isArgumentIndexOptional)
       throws PiranhaConfigurationException {
-    String methodName = getValueStringFromMap(methodPropertyEntry, METHOD_NAME_KEY);
-    String flagType = getValueStringFromMap(methodPropertyEntry, FLAG_TYPE_KEY);
     Integer argumentIndexInteger = getArgumentIndexFromMap(methodPropertyEntry, ARGUMENT_INDEX_KEY);
-    if (methodName == null) {
+    if (!isArgumentIndexOptional && argumentIndexInteger == null) {
       throw new PiranhaConfigurationException(
-          "methodProperty is missing mandatory methodName field. Check:\n" + methodPropertyEntry);
-    } else if (flagType == null) {
-      throw new PiranhaConfigurationException(
-          "methodProperty is missing mandatory flagType field. Check:\n" + methodPropertyEntry);
-    } else if (!isArgumentIndexOptional && argumentIndexInteger == null) {
-      throw new PiranhaConfigurationException(
-          "methodProperty did not have argumentIndex. By default, Piranha requires an argument index for flag "
+          propertyName
+              + " did not have argumentIndex. By default, Piranha requires an argument index for flag "
               + "APIs, to which the flag name/symbol will be passed. This is to avoid over-deletion of all "
               + "occurrences of a flag API method. If you are sure you want to delete all instances of the "
               + "method below, consider using Piranha:ArgumentIndexOptional=true to override this behavior. "
               + "Check:\n"
               + methodPropertyEntry);
-    } else if (argumentIndexInteger != null && argumentIndexInteger.intValue() < 0) {
+    } else if (argumentIndexInteger != null && argumentIndexInteger < 0) {
       throw new PiranhaConfigurationException(
           "Invalid argumentIndex field. Arguments are zero indexed. Check:\n"
               + methodPropertyEntry);
     }
-
-    return new PiranhaMethodRecord(
-        methodName,
-        flagType,
-        Optional.ofNullable(argumentIndexInteger),
-        Optional.ofNullable(getValueStringFromMap(methodPropertyEntry, RECEIVER_TYPE_STRING)),
-        Optional.ofNullable(getValueStringFromMap(methodPropertyEntry, RETURN_TYPE_STRING)));
+    return argumentIndexInteger;
   }
 }
diff --git a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaRecord.java b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaRecord.java
index 55b5d90a0..62252e43c 100644
--- a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaRecord.java
+++ b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaRecord.java
@@ -16,7 +16,10 @@
 import java.util.Map;
 import javax.annotation.Nullable;
 
-/** Utility methods used by {@link PiranhaEnumRecord} and {@link PiranhaMethodRecord} classes */
+/**
+ * Utility methods used by {@link PiranhaEnumRecord}, {@link PiranhaFlagMethodRecord}, and {@link
+ * PiranhaTestMethodRecord} classes
+ */
 public class PiranhaRecord {
   /**
    * Utility method. Checks whether the value associated to a given map and given key is a non-empty
diff --git a/java/piranha/src/main/java/com/uber/piranha/config/PiranhaTestMethodRecord.java b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaTestMethodRecord.java
new file mode 100644
index 000000000..fb096813e
--- /dev/null
+++ b/java/piranha/src/main/java/com/uber/piranha/config/PiranhaTestMethodRecord.java
@@ -0,0 +1,90 @@
+/**
+ * Copyright (c) 2021 Uber Technologies, Inc.
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.uber.piranha.config;
+
+import static com.uber.piranha.config.PiranhaMethodRecord.parseArgumentIndexFromJSON;
+import static com.uber.piranha.config.PiranhaMethodRecord.parseMethodNameFromJSON;
+import static com.uber.piranha.config.PiranhaRecord.getValueStringFromMap;
+
+import java.util.Map;
+import java.util.Optional;
+
+/** A class representing a test method configuration record from properties.json */
+public final class PiranhaTestMethodRecord implements PiranhaMethodRecord {
+  // Used for generating exception messages
+  private static final String PROPERTY_NAME = "testMethodProperty";
+
+  private final String methodName;
+  private final Optional<Integer> argumentIdx;
+  private final Optional<String> receiverType;
+  private final Optional<String> returnType;
+
+  PiranhaTestMethodRecord(
+      String methodName,
+      Optional<Integer> argumentIdx,
+      Optional<String> receiverType,
+      Optional<String> returnType) {
+    this.methodName = methodName;
+    this.argumentIdx = argumentIdx;
+    this.receiverType = receiverType;
+    this.returnType = returnType;
+  }
+
+  @Override
+  public String getMethodName() {
+    return methodName;
+  }
+
+  @Override
+  public Optional<Integer> getArgumentIdx() {
+    return argumentIdx;
+  }
+
+  @Override
+  public Optional<String> getReceiverType() {
+    return receiverType;
+  }
+
+  @Override
+  public Optional<String> getReturnType() {
+    return returnType;
+  }
+
+  /**
+   * Parse the entry for a single method from piranha.json that has been previously decoded into a
+   * map
+   *
+   * @param methodPropertyEntry The decoded json entry (as a Map of property names to values)
+   * @param isArgumentIndexOptional Whether argumentIdx should be treated as optional
+   * @return A PiranhaTestMethodRecord corresponding to the given map/json record.
+   * @throws PiranhaConfigurationException if there was any issue reading or parsing the
+   *     configuration file.
+   */
+  static PiranhaTestMethodRecord parseFromJSONPropertyEntryMap(
+      Map<String, Object> methodPropertyEntry, boolean isArgumentIndexOptional)
+      throws PiranhaConfigurationException {
+
+    String methodName = parseMethodNameFromJSON(PROPERTY_NAME, methodPropertyEntry);
+    Integer argumentIndexInteger =
+        parseArgumentIndexFromJSON(PROPERTY_NAME, methodPropertyEntry, isArgumentIndexOptional);
+    String receiverType = getValueStringFromMap(methodPropertyEntry, RECEIVER_TYPE_STRING);
+    String returnType = getValueStringFromMap(methodPropertyEntry, RETURN_TYPE_STRING);
+
+    return new PiranhaTestMethodRecord(
+        methodName,
+        Optional.ofNullable(argumentIndexInteger),
+        Optional.ofNullable(receiverType),
+        Optional.ofNullable(returnType));
+  }
+}
diff --git a/java/piranha/src/test/java/com/uber/piranha/TestCaseCleanUpTest.java b/java/piranha/src/test/java/com/uber/piranha/TestCaseCleanUpTest.java
index b19028342..2e24fae3c 100644
--- a/java/piranha/src/test/java/com/uber/piranha/TestCaseCleanUpTest.java
+++ b/java/piranha/src/test/java/com/uber/piranha/TestCaseCleanUpTest.java
@@ -833,4 +833,196 @@ public void testCleanBySettersHeuristicIgnoreSettersForOtherFlags() throws IOExc
             "}")
         .doTest();
   }
+
+  /**
+   * Tests the removal of extraneous test methods after a feature flag is removed. An example of
+   * this might be an attempt to mock the result of a feature flag call such as
+   * mock(experimentation.isToggleEnabled(STALE_FLAG));
+   */
+  @Test
+  public void testTestMethodRemoval() throws IOException {
+    ErrorProneFlags.Builder bFlag = ErrorProneFlags.builder();
+    bFlag.putFlag("Piranha:FlagName", "STALE_FLAG");
+    bFlag.putFlag("Piranha:IsTreated", "false");
+    bFlag.putFlag("Piranha:Config", "config/properties.json");
+
+    BugCheckerRefactoringTestHelper bcr =
+        BugCheckerRefactoringTestHelper.newInstance(new XPFlagCleaner(bFlag.build()), getClass());
+
+    bcr = bcr.setArgs("-d", temporaryFolder.getRoot().getAbsolutePath());
+
+    bcr = PiranhaTestingHelpers.addHelperClasses(bcr);
+    bcr = addExperimentFlagEnums(bcr); // Adds STALE_FLAG, etc enums
+    bcr.addInputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Before;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestExperimentName.STALE_FLAG;",
+            "import static com.uber.piranha.TestMethods.mock;",
+            "import static com.uber.piranha.TestMethods.mockable;",
+            "import static com.uber.piranha.TestMethods.keepMe;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Before",
+            "  public void setUp() {",
+            "     mock(experimentation.isToggleDisabled(STALE_FLAG)).thenReturn(false);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG));",
+            "     mock(mockable());",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "     keepMe(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag() {",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "  }",
+            "}")
+        .addOutputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Before;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestMethods.mock;",
+            "import static com.uber.piranha.TestMethods.mockable;",
+            "import static com.uber.piranha.TestMethods.keepMe;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Before",
+            "  public void setUp() {",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     mock(mockable());",
+            "     keepMe(false).thenReturn(true);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag() {",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  /**
+   * Unfortunately, test method removal isn't smart enough to clean up imports without requiring a
+   * second pass. This test ensures the imports still exist. However, there are tools that exist
+   * that can be run to remove unused imports.
+   */
+  @Test
+  public void testTestMethodRemovalKeepImports() throws IOException {
+    ErrorProneFlags.Builder bFlag = ErrorProneFlags.builder();
+    bFlag.putFlag("Piranha:FlagName", "STALE_FLAG");
+    bFlag.putFlag("Piranha:IsTreated", "false");
+    bFlag.putFlag("Piranha:Config", "config/properties.json");
+
+    BugCheckerRefactoringTestHelper bcr =
+        BugCheckerRefactoringTestHelper.newInstance(new XPFlagCleaner(bFlag.build()), getClass());
+
+    bcr = bcr.setArgs("-d", temporaryFolder.getRoot().getAbsolutePath());
+
+    bcr = PiranhaTestingHelpers.addHelperClasses(bcr);
+    bcr = addExperimentFlagEnums(bcr); // Adds STALE_FLAG, etc enums
+    bcr.addInputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Before;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestExperimentName.STALE_FLAG;",
+            "import static com.uber.piranha.TestMethods.mock;",
+            "import static com.uber.piranha.TestMethods.keepMe;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Before",
+            "  public void setUp() {",
+            "     mock(experimentation.isToggleDisabled(STALE_FLAG)).thenReturn(false);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG));",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "     keepMe(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag() {",
+            "     mock(experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "  }",
+            "}")
+        .addOutputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Before;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestMethods.mock;",
+            "import static com.uber.piranha.TestMethods.keepMe;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Before",
+            "  public void setUp() {",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     keepMe(false).thenReturn(true);",
+            "  }",
+            "  @Test",
+            "  public void test_StaleFlag() {",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  /**
+   * A test when specifying a particular index when cleaning a test method. This test will ensure
+   * that only methods that have the flag method in the right index AND that that flag method
+   * matches a flag will be cleaned up.
+   */
+  @Test
+  public void testTestOverloadedMethodRemoval() throws IOException {
+    ErrorProneFlags.Builder bFlag = ErrorProneFlags.builder();
+    bFlag.putFlag("Piranha:FlagName", "STALE_FLAG");
+    bFlag.putFlag("Piranha:IsTreated", "false");
+    bFlag.putFlag("Piranha:Config", "config/properties.json");
+
+    BugCheckerRefactoringTestHelper bcr =
+        BugCheckerRefactoringTestHelper.newInstance(new XPFlagCleaner(bFlag.build()), getClass());
+
+    bcr = bcr.setArgs("-d", temporaryFolder.getRoot().getAbsolutePath());
+
+    bcr = PiranhaTestingHelpers.addHelperClasses(bcr);
+    bcr = addExperimentFlagEnums(bcr); // Adds STALE_FLAG, etc enums
+    bcr.addInputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestExperimentName.OTHER_FLAG_1;",
+            "import static com.uber.piranha.TestExperimentName.STALE_FLAG;",
+            "import static com.uber.piranha.TestMethods.expect;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     expect(experimentation.isToggleEnabled(STALE_FLAG));",
+            "     expect(true, experimentation.isToggleEnabled(STALE_FLAG)).thenReturn(true);",
+            "     expect(true, experimentation.isToggleEnabled(OTHER_FLAG_1)).thenReturn(true);",
+            "  }",
+            "}")
+        .addOutputLines(
+            "TestClass.java",
+            "package com.uber.piranha;",
+            "import org.junit.Test;",
+            "import static com.uber.piranha.TestExperimentName.OTHER_FLAG_1;",
+            "import static com.uber.piranha.TestMethods.expect;",
+            "class TestClass {",
+            "  private XPTest experimentation;",
+            "  @Test",
+            "  public void test_StaleFlag1() {",
+            "     expect(false);",
+            "     expect(true, experimentation.isToggleEnabled(OTHER_FLAG_1)).thenReturn(true);",
+            "  }",
+            "}")
+        .doTest();
+  }
 }
diff --git a/java/piranha/src/test/resources/com/uber/piranha/TestMethods.java b/java/piranha/src/test/resources/com/uber/piranha/TestMethods.java
new file mode 100644
index 000000000..2c85e3fc9
--- /dev/null
+++ b/java/piranha/src/test/resources/com/uber/piranha/TestMethods.java
@@ -0,0 +1,42 @@
+/**
+ * Copyright (c) 2021 Uber Technologies, Inc.
+ *
+ * <p>Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file
+ * except in compliance with the License. You may obtain a copy of the License at
+ *
+ * <p>http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * <p>Unless required by applicable law or agreed to in writing, software distributed under the
+ * License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.uber.piranha;
+
+class MockObject {
+  public boolean thenReturn(boolean bool) {
+    return bool;
+  }
+}
+
+class TestMethods {
+  public static boolean mockable() {
+    return false;
+  }
+
+  public static MockObject mock(Object thingToMock) {
+    return new MockObject();
+  }
+
+  public static MockObject keepMe(Object thingToMock) {
+    return new MockObject();
+  }
+
+  public static MockObject expect(Object thingToMock) {
+    return new MockObject();
+  }
+
+  public static MockObject expect(boolean value, Object thingToMock) {
+    return new MockObject();
+  }
+}