fix tests

google-gemini · Feb 27, 2024 · 03cebef · 03cebef
1 parent 6f60cc7
commit 03cebef
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 2 deletions.
diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/StreamingSnapshotTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/StreamingSnapshotTests.kt
@@ -19,6 +19,7 @@ package com.google.ai.client.generativeai
 import com.google.ai.client.generativeai.type.BlockReason
 import com.google.ai.client.generativeai.type.FinishReason
 import com.google.ai.client.generativeai.type.HarmCategory
+import com.google.ai.client.generativeai.type.InvalidAPIKeyException
 import com.google.ai.client.generativeai.type.PromptBlockedException
 import com.google.ai.client.generativeai.type.ResponseStoppedException
 import com.google.ai.client.generativeai.type.SerializationException
@@ -173,6 +174,6 @@ internal class StreamingSnapshotTests {
     goldenStreamingFile("failure-api-key.txt", HttpStatusCode.BadRequest) {
       val responses = model.generateContentStream()
 
-      withTimeout(testTimeout) { shouldThrow<ServerException> { responses.collect() } }
+      withTimeout(testTimeout) { shouldThrow<InvalidAPIKeyException> { responses.collect() } }
     }
 }
diff --git a/generativeai/src/test/java/com/google/ai/client/generativeai/UnarySnapshotTests.kt b/generativeai/src/test/java/com/google/ai/client/generativeai/UnarySnapshotTests.kt
@@ -19,6 +19,7 @@ package com.google.ai.client.generativeai
 import com.google.ai.client.generativeai.type.BlockReason
 import com.google.ai.client.generativeai.type.FinishReason
 import com.google.ai.client.generativeai.type.HarmCategory
+import com.google.ai.client.generativeai.type.InvalidAPIKeyException
 import com.google.ai.client.generativeai.type.PromptBlockedException
 import com.google.ai.client.generativeai.type.ResponseStoppedException
 import com.google.ai.client.generativeai.type.SerializationException
@@ -129,7 +130,7 @@ internal class UnarySnapshotTests {
   @Test
   fun `invalid api key`() =
     goldenUnaryFile("failure-api-key.json", HttpStatusCode.BadRequest) {
-      withTimeout(testTimeout) { shouldThrow<ServerException> { model.generateContent() } }
+      withTimeout(testTimeout) { shouldThrow<InvalidAPIKeyException> { model.generateContent() } }
     }
 
   @Test