Add support for AssertJ as() and describedAs() in AssertionHandler

nullaway/src/main/java/com/uber/nullaway/handlers/AssertionHandler.java

@@ -30,6 +30,7 @@
 import com.uber.nullaway.dataflow.AccessPath;
 import com.uber.nullaway.dataflow.AccessPathNullnessPropagation;
 import java.util.List;
+import javax.annotation.Nullable;
 import org.checkerframework.nullaway.dataflow.cfg.node.MethodInvocationNode;
 import org.checkerframework.nullaway.dataflow.cfg.node.Node;
 
@@ -60,17 +61,9 @@ public NullnessHint onDataflowVisitMethodInvocation(
     // assertThat(A).isInstanceOf(Foo.class)
     // A will not be NULL after this statement.
     if (methodNameUtil.isMethodIsNotNull(callee) || methodNameUtil.isMethodIsInstanceOf(callee)) {
-      Node receiver = node.getTarget().getReceiver();
-      if (receiver instanceof MethodInvocationNode) {
-        MethodInvocationNode receiver_method = (MethodInvocationNode) receiver;
-        Symbol.MethodSymbol receiver_symbol = ASTHelpers.getSymbol(receiver_method.getTree());
-        if (methodNameUtil.isMethodAssertThat(receiver_symbol)) {
-          Node arg = receiver_method.getArgument(0);
-          AccessPath ap = AccessPath.getAccessPathForNode(arg, state, apContext);
-          if (ap != null) {
-            bothUpdates.set(ap, NONNULL);
-          }
-        }
+      AccessPath ap = getAccessPathForNotNullAssertThatExpr(node, state, apContext);
+      if (ap != null) {
+        bothUpdates.set(ap, NONNULL);
       }
     }
 
@@ -94,4 +87,31 @@ public NullnessHint onDataflowVisitMethodInvocation(
 
     return NullnessHint.UNKNOWN;
   }
+
+  /**
+   * Returns the AccessPath for the argument of an assertThat() call, if present as a valid nested
+   * receiver expression of a method invocation
+   *
+   * @param node the method invocation node
+   * @param state the visitor state
+   * @param apContext the access path context
+   * @return the AccessPath for the argument of the assertThat() call, if present, otherwise {@code
+   *     null}
+   */
+  private @Nullable AccessPath getAccessPathForNotNullAssertThatExpr(
+      MethodInvocationNode node, VisitorState state, AccessPath.AccessPathContext apContext) {
+    Node receiver = node.getTarget().getReceiver();
+    if (receiver instanceof MethodInvocationNode) {
+      MethodInvocationNode receiver_method = (MethodInvocationNode) receiver;
+      Symbol.MethodSymbol receiver_symbol = ASTHelpers.getSymbol(receiver_method.getTree());
+      if (methodNameUtil.isMethodAssertThat(receiver_symbol)) {
+        Node arg = receiver_method.getArgument(0);
+        return AccessPath.getAccessPathForNode(arg, state, apContext);
+      } else if (methodNameUtil.isMethodAssertJDescribedAs(receiver_symbol)) {
+        // For calls to as() or describedAs(), we recursively search for the assertThat() call
+        return getAccessPathForNotNullAssertThatExpr(receiver_method, state, apContext);
+      }
+    }
+    return null;
+  }
 }

nullaway/src/main/java/com/uber/nullaway/handlers/MethodNameUtil.java

@@ -56,6 +56,9 @@ class MethodNameUtil {
   private static final String IS_PRESENT_OWNER_ASSERTJ =
       "org.assertj.core.api.AbstractOptionalAssert";
   private static final String ASSERT_THAT_METHOD = "assertThat";
+  private static final String AS_METHOD = "as";
+  private static final String DESCRIBED_AS_METHOD = "describedAs";
+
   private static final String ASSERT_THAT_OWNER_TRUTH = "com.google.common.truth.Truth";
   private static final String ASSERT_THAT_OWNER_ASSERTJ = "org.assertj.core.api.Assertions";
 
@@ -101,6 +104,9 @@ class MethodNameUtil {
   private Name assertThatOwnerTruth;
   private Name assertThatOwnerAssertJ;
 
+  private Name as;
+  private Name describedAs;
+
   // Names for junit assertion libraries.
   private Name hamcrestAssertClass;
   private Name junitAssertClass;
@@ -141,6 +147,9 @@ void initializeMethodNames(Name.Table table) {
     assertThatOwnerTruth = table.fromString(ASSERT_THAT_OWNER_TRUTH);
     assertThatOwnerAssertJ = table.fromString(ASSERT_THAT_OWNER_ASSERTJ);
 
+    as = table.fromString(AS_METHOD);
+    describedAs = table.fromString(DESCRIBED_AS_METHOD);
+
     isPresent = table.fromString(IS_PRESENT_METHOD);
     isNotEmpty = table.fromString(IS_NOT_EMPTY_METHOD);
     isPresentOwnerAssertJ = table.fromString(IS_PRESENT_OWNER_ASSERTJ);
@@ -211,6 +220,18 @@ boolean isMethodAssertThat(Symbol.MethodSymbol methodSymbol) {
         || matchesMethod(methodSymbol, assertThat, assertThatOwnerAssertJ);
   }
 
+  /**
+   * Returns true if the method is describedAs() or as() from AssertJ. Note that this implementation
+   * does not check the ower, as there are many possible implementations. This method should only be
+   * used in a caller content where it is clear that the operation is related to use of AssertJ.
+   *
+   * @param methodSymbol symbol for the method
+   * @return {@code true} iff the method is describedAs() or as() from AssertJ
+   */
+  public boolean isMethodAssertJDescribedAs(Symbol.MethodSymbol methodSymbol) {
+    return methodSymbol.name.equals(as) || methodSymbol.name.equals(describedAs);
+  }
+
   boolean isMethodHamcrestAssertThat(Symbol.MethodSymbol methodSymbol) {
     return matchesMethod(methodSymbol, assertThat, hamcrestAssertClass);
   }

nullaway/src/test/java/com/uber/nullaway/NullAwayAssertionLibsTests.java

@@ -365,6 +365,71 @@ public void supportAssertJAssertThatIsNotNull_Object() {
         .doTest();
   }
 
+  @Test
+  public void supportAssertJAssertThatIsNotNullWithDescription_Object() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:HandleTestAssertionLibraries=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.lang.Object;",
+            "import java.util.Objects;",
+            "import javax.annotation.Nullable;",
+            "import static org.assertj.core.api.Assertions.assertThat;",
+            "class Test {",
+            "  private void foo(@Nullable Object o) {",
+            "    assertThat(o).as(\"test\").isNotNull();",
+            "    o.toString();",
+            "  }",
+            "  private void foo2(@Nullable Object o) {",
+            "    assertThat(o).describedAs(\"test\").isNotNull();",
+            "    o.toString();",
+            "  }",
+            "  private void foo3(@Nullable Object o) {",
+            "    assertThat(o).describedAs(\"test1\").as(\"test2\").isNotNull();",
+            "    o.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
+  @Test
+  public void assertJAssertThatIsNotNullUnhandled() {
+    makeTestHelperWithArgs(
+            Arrays.asList(
+                "-d",
+                temporaryFolder.getRoot().getAbsolutePath(),
+                "-XepOpt:NullAway:AnnotatedPackages=com.uber",
+                "-XepOpt:NullAway:HandleTestAssertionLibraries=true"))
+        .addSourceLines(
+            "Test.java",
+            "package com.uber;",
+            "import java.lang.Object;",
+            "import java.util.Objects;",
+            "import javax.annotation.Nullable;",
+            "import static org.assertj.core.api.Assertions.assertThat;",
+            "class Test {",
+            "  private void foo(@Nullable Object o) {",
+            "    org.assertj.core.api.ObjectAssert t = assertThat(o);",
+            "    t.isNotNull();",
+            "    // False positive",
+            "    // BUG: Diagnostic contains: dereferenced expression",
+            "    o.toString();",
+            "  }",
+            "  private void foo2(@Nullable Object o) {",
+            "    assertThat(o).isEqualToIgnoringNullFields(o).describedAs(\"test\").isNotNull();",
+            "    // False positive",
+            "    // BUG: Diagnostic contains: dereferenced expression",
+            "    o.toString();",
+            "  }",
+            "}")
+        .doTest();
+  }
+
   @Test
   public void supportAssertJAssertThatIsNotNull_String() {
     makeTestHelperWithArgs(