fix streaming unknown enum test

google-gemini · Jul 15, 2024 · c1fbf0b · c1fbf0b
1 parent 5c488b1
commit c1fbf0b
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/common/src/test/java/com/google/ai/client/generativeai/common/StreamingSnapshotTests.kt b/common/src/test/java/com/google/ai/client/generativeai/common/StreamingSnapshotTests.kt
@@ -72,11 +72,13 @@ internal class StreamingSnapshotTests {
       val responses = apiController.generateContentStream(textGenerateContentRequest("prompt"))
 
       withTimeout(testTimeout) {
-        responses.first {
+        val responseList = responses.toList()
+        responseList.isEmpty() shouldBe false
+        responseList.any {
           it.candidates?.any {
             it.safetyRatings?.any { it.category == HarmCategory.UNKNOWN } ?: false
           } ?: false
-        }
+        } shouldBe true
       }
     }