Implement return type matches and do not flag variable class arguments

error-prone-contrib/src/main/java/tech/picnic/errorprone/bugpatterns/MockitoMockClassReference.java

@@ -8,8 +8,12 @@
 import static com.google.errorprone.matchers.Matchers.argument;
 import static com.google.errorprone.matchers.Matchers.argumentCount;
 import static com.google.errorprone.matchers.Matchers.isSameType;
+import static com.google.errorprone.matchers.Matchers.isVariable;
+import static com.google.errorprone.matchers.Matchers.not;
 import static com.google.errorprone.matchers.Matchers.staticMethod;
+import static com.sun.source.tree.Tree.Kind.RETURN;
 import static com.sun.source.tree.Tree.Kind.VARIABLE;
+import static java.util.Objects.requireNonNull;
 import static tech.picnic.errorprone.bugpatterns.util.Documentation.BUG_PATTERNS_BASE_URL;
 
 import com.google.auto.service.AutoService;
@@ -23,6 +27,8 @@
 import com.google.errorprone.matchers.Matcher;
 import com.google.errorprone.util.ASTHelpers;
 import com.sun.source.tree.MethodInvocationTree;
+import com.sun.source.tree.MethodTree;
+import com.sun.source.tree.ReturnTree;
 import com.sun.source.tree.Tree;
 import com.sun.source.tree.VariableTree;
 
@@ -45,11 +51,13 @@ public final class MockitoMockClassReference extends BugChecker
   private static final Matcher<MethodInvocationTree> MOCKITO_MOCK_OR_SPY =
       allOf(
           argumentCount(1),
-          argument(0, isSameType(Class.class.getName())),
+          argument(0, allOf(isSameType(Class.class.getName()), not(isVariable()))),
           staticMethod().onClass("org.mockito.Mockito").namedAnyOf("mock", "spy"));
   // XXX: Replace `var` usage with explicit type instead.
   private static final Matcher<VariableTree> INCOMPATIBLE_VARIABLE =
       anyOf(ASTHelpers::hasNoExplicitType, MockitoMockClassReference::hasTypeDifference);
+  private static final Matcher<ReturnTree> INCOMPATIBLE_RETURN =
+      MockitoMockClassReference::hasTypeDifference;
 
   /** Instantiates a new {@link MockitoMockClassReference} instance. */
   public MockitoMockClassReference() {}
@@ -60,18 +68,28 @@ public Description matchMethodInvocation(MethodInvocationTree tree, VisitorState
       return Description.NO_MATCH;
     }
 
-    // XXX: Add similar matchers for usage in a ReturnTree and the method's return type.
     Tree parent = state.getPath().getParentPath().getLeaf();
     if (parent.getKind() == VARIABLE
         && INCOMPATIBLE_VARIABLE.matches((VariableTree) parent, state)) {
       return Description.NO_MATCH;
+    } else if (parent.getKind() == RETURN
+        && INCOMPATIBLE_RETURN.matches((ReturnTree) parent, state)) {
+      return Description.NO_MATCH;
     }
 
     return describeMatch(tree, SuggestedFix.delete(Iterables.getOnlyElement(tree.getArguments())));
   }
 
   private static boolean hasTypeDifference(VariableTree tree, VisitorState state) {
-    return !ASTHelpers.isSameType(
-        ASTHelpers.getType(tree), ASTHelpers.getType(tree.getInitializer()), state);
+    return hasTypeDifference(tree, tree.getInitializer(), state);
+  }
+
+  private static boolean hasTypeDifference(ReturnTree tree, VisitorState state) {
+    Tree returnTypeTree = requireNonNull(state.findEnclosing(MethodTree.class)).getReturnType();
+    return hasTypeDifference(returnTypeTree, tree.getExpression(), state);
+  }
+
+  private static boolean hasTypeDifference(Tree treeA, Tree treeB, VisitorState state) {
+    return !ASTHelpers.isSameType(ASTHelpers.getType(treeA), ASTHelpers.getType(treeB), state);
   }
 }

error-prone-contrib/src/test/java/tech/picnic/errorprone/bugpatterns/MockitoMockClassReferenceTest.java

@@ -28,49 +28,37 @@ void identification() {
             "    variableMock = mock(Integer.class);",
             "    // BUG: Diagnostic contains:",
             "    List nonGenericallyTypedMock = mock(List.class);",
-            "    var dynamicallyTypedMock = mock(Integer.class);",
+            "    // BUG: Diagnostic contains:",
+            "    List<String> genericallyTypedMock = mock(List.class);",
+            "    var nonExplicitlyTypedMock = mock(Integer.class);",
+            "    Class<? extends Number> variableType = Integer.class;",
+            "    Number variableTypedMock = mock(variableType);",
             "    Integer namedMock = mock(Integer.class, \"name\");",
             "    Object subtypeMock = mock(Integer.class);",
             "",
             "    // BUG: Diagnostic contains:",
             "    String equalTypedSpy = spy(String.class);",
             "    // BUG: Diagnostic contains:",
             "    List nonGenericallyTypedSpy = spy(List.class);",
-            "    var dynamicallyTypedSpy = spy(Integer.class);",
+            "    // BUG: Diagnostic contains:",
+            "    List<String> genericallyTypedSpy = spy(List.class);",
+            "    var nonExplicitlyTypedSpy = spy(Integer.class);",
+            "    Number subtypeSpy = spy(Integer.class);",
             "    Object objectSpy = spy(new Object());",
             "  }",
-            "}")
-        .doTest();
-  }
-
-  @Test
-  void unimplementedCases() {
-    // XXX: Move to identification once fixed.
-    CompilationTestHelper.newInstance(MockitoMockClassReference.class, getClass())
-        .addSourceLines(
-            "A.java",
-            "import static org.mockito.Mockito.mock;",
             "",
-            "import java.math.BigInteger;",
-            "import java.util.List;",
-            "",
-            "class A {",
-            "  void m() {",
-            "    // This case is arguable as it is unsafe in any case.",
-            "    // This currently is not identified as the `erasure` is the same.",
-            "    List<String> genericallyTypedMock = mock(List.class);",
-            "",
-            "    // This case is problematic and should not be replaced.",
-            "    // The real type of `variableTypedMock` is `BigInteger`.",
-            "    // But when replaced, it will be `Number`.",
-            "    Class<? extends Number> variableType = BigInteger.class;",
-            "    Number variableTypedMock = mock(variableType);",
+            "  Integer getIntegerMock() {",
+            "    // BUG: Diagnostic contains:",
+            "    return mock(Integer.class);",
             "  }",
             "",
-            "  <T extends Number> T getGenericMock(Class<T> clazz) {",
-            "    // No idea how to detect type T as part of generics usage.",
+            "  <T> T getGenericMock(Class<T> clazz) {",
             "    return mock(clazz);",
             "  }",
+            "",
+            "  Number getSubTypeMock() {",
+            "    return mock(Integer.class);",
+            "  }",
             "}")
         .doTest();
   }
@@ -104,6 +92,10 @@ void replacement() {
             "    var unknownTypeSpy = spy(Integer.class);",
             "    Object objectSpy = spy(new Object());",
             "  }",
+            "",
+            "  Integer getIntegerMock() {",
+            "    return mock(Integer.class);",
+            "  }",
             "}")
         .addOutputLines(
             "A.java",
@@ -131,6 +123,10 @@ void replacement() {
             "    var unknownTypeSpy = spy(Integer.class);",
             "    Object objectSpy = spy(new Object());",
             "  }",
+            "",
+            "  Integer getIntegerMock() {",
+            "    return mock();",
+            "  }",
             "}")
         .doTest(TestMode.TEXT_MATCH);
   }