diff --git a/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt b/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt
index ad5c634a..dcb37a93 100644
--- a/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt
+++ b/common/src/test/java/com/google/ai/client/generativeai/common/UnarySnapshotTests.kt
@@ -80,9 +80,9 @@ internal class UnarySnapshotTests {
       withTimeout(testTimeout) {
         val response = apiController.generateContent(textGenerateContentRequest("prompt"))
 
-        response.candidates?.first {
-          it.safetyRatings?.any { it.category == HarmCategory.UNKNOWN } ?: false
-        }
+        response.candidates?.isNullOrEmpty() shouldBe false
+        val candidate = response.candidates?.first()
+        candidate?.safetyRatings?.any { it.category == HarmCategory.UNKNOWN } shouldBe true
       }
     }