From c1fbf0b53170048effbafd1f85628edd05f57bd3 Mon Sep 17 00:00:00 2001
From: Tanzim Hossain <tanzim@google.com>
Date: Fri, 12 Jul 2024 15:42:59 -0400
Subject: [PATCH] fix streaming unknown enum test

---
 .../ai/client/generativeai/common/StreamingSnapshotTests.kt | 6 ++++--
 1 file changed, 4 insertions(+), 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
index 7f151320..656a5eb4 100644
--- 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
       }
     }