diff --git a/java/tests/instrumentation/src/com/android/textclassifier/subjects/EntitySubject.java b/java/tests/instrumentation/src/com/android/textclassifier/subjects/EntitySubject.java
index 6d01a64..30a9dca 100644
--- a/java/tests/instrumentation/src/com/android/textclassifier/subjects/EntitySubject.java
+++ b/java/tests/instrumentation/src/com/android/textclassifier/subjects/EntitySubject.java
@@ -20,7 +20,6 @@
 
 import com.android.textclassifier.Entity;
 import com.google.common.truth.FailureMetadata;
-import com.google.common.truth.MathUtil;
 import com.google.common.truth.Subject;
 import javax.annotation.Nullable;
 
@@ -44,8 +43,7 @@
     if (!entity.getEntityType().equals(this.entity.getEntityType())) {
       failWithActual("expected to have type", entity.getEntityType());
     }
-    if (!MathUtil.equalWithinTolerance(entity.getScore(), this.entity.getScore(), TOLERANCE)) {
-      failWithActual("expected to have confidence score", entity.getScore());
-    }
+    check("expected to have confidence score").that(entity.getScore()).isWithin(TOLERANCE)
+      .of(this.entity.getScore());
   }
 }