diff --git a/nullaway/src/main/java/com/uber/nullaway/NullAway.java b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
index 4104da418d..3d17824579 100644
--- a/nullaway/src/main/java/com/uber/nullaway/NullAway.java
+++ b/nullaway/src/main/java/com/uber/nullaway/NullAway.java
@@ -639,12 +639,14 @@ private Description checkReturnExpression(
       return Description.NO_MATCH;
     }
     if (mayBeNullExpr(state, retExpr)) {
-      String message = "returning @Nullable expression from method with @NonNull return type";
+      final ErrorMessage errorMessage =
+          new ErrorMessage(
+              MessageTypes.RETURN_NULLABLE,
+              "returning @Nullable expression from method with @NonNull return type");
+      handler.onPrepareErrorMessage(retExpr, state, errorMessage);
+
       return errorBuilder.createErrorDescriptionForNullAssignment(
-          new ErrorMessage(MessageTypes.RETURN_NULLABLE, message),
-          retExpr,
-          state.getPath(),
-          buildDescription(tree));
+          errorMessage, retExpr, state.getPath(), buildDescription(tree));
     }
     return Description.NO_MATCH;
   }
diff --git a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
index 41cb0dc545..3010628a26 100644
--- a/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
+++ b/nullaway/src/test/java/com/uber/nullaway/NullAwayTest.java
@@ -1575,7 +1575,7 @@ public void OptionalEmptinessRxPositiveTest() {
             "           .map(optional -> optional.get().toString());",
             "     observable",
             "           .filter(optional -> optional.isPresent() || perhaps())",
-            "           // BUG: Diagnostic contains: returning @Nullable expression from method with @NonNull",
+            "           // BUG: Diagnostic contains: Optional optional can be empty",
             "           .map(optional -> optional.get())",
             "           .map(irr -> irr.toString());",
             "     }",
@@ -1602,6 +1602,9 @@ public void OptionalEmptinessRxNegativeTest() {
             "  private static boolean perhaps() { return Math.random() > 0.5; }",
             "  void foo(Observable<Optional<String>> observable) {",
             "     observable",
+            "           .filter(optional -> optional.isPresent())",
+            "           .map(optional -> optional.get().toString());",
+            "     observable",
             "           .filter(optional -> optional.isPresent() && perhaps())",
             "           .map(optional -> optional.get().toString());",
             "     observable",