Skip to content

Commit 09a7c24

Browse files
Add tests for HarmBlockMethod compatibility with GoogleAI and VertexAI
Added unit tests to `GenerativeModelTesting.kt` to verify that: 1. Using `HarmBlockMethod` with `GoogleAI` backend throws `InvalidStateException`. 2. Using `HarmBlockMethod` with `VertexAI` backend does not throw `InvalidStateException`. This covers an edge case where `HarmBlockMethod` is only supported by VertexAI.
1 parent e4369ab commit 09a7c24

File tree

1 file changed

+97
-0
lines changed

1 file changed

+97
-0
lines changed

firebase-ai/src/test/java/com/google/firebase/ai/GenerativeModelTesting.kt

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,14 @@ import com.google.firebase.ai.common.util.doBlocking
2323
import com.google.firebase.ai.type.Candidate
2424
import com.google.firebase.ai.type.Content
2525
import com.google.firebase.ai.type.GenerateContentResponse
26+
import com.google.firebase.ai.type.GenerativeBackend
27+
import com.google.firebase.ai.type.HarmBlockMethod
28+
import com.google.firebase.ai.type.HarmBlockThreshold
29+
import com.google.firebase.ai.type.HarmCategory
30+
import com.google.firebase.ai.type.InvalidStateException
2631
import com.google.firebase.ai.type.PublicPreviewAPI
2732
import com.google.firebase.ai.type.RequestOptions
33+
import com.google.firebase.ai.type.SafetySetting
2834
import com.google.firebase.ai.type.ServerException
2935
import com.google.firebase.ai.type.TextPart
3036
import com.google.firebase.ai.type.content
@@ -146,6 +152,97 @@ internal class GenerativeModelTesting {
146152
exception.message shouldContain "location"
147153
}
148154

155+
@Test
156+
fun `exception thrown when using HarmBlockMethod with GoogleAI`() = doBlocking {
157+
val mockEngine = MockEngine {
158+
respond(
159+
generateContentResponseAsJsonString("text response"),
160+
HttpStatusCode.OK,
161+
headersOf(HttpHeaders.ContentType, "application/json")
162+
)
163+
}
164+
165+
val apiController =
166+
APIController(
167+
"super_cool_test_key",
168+
"gemini-2.5-flash",
169+
RequestOptions(),
170+
mockEngine,
171+
TEST_CLIENT_ID,
172+
mockFirebaseApp,
173+
TEST_VERSION,
174+
TEST_APP_ID,
175+
null,
176+
)
177+
178+
val safetySettings =
179+
listOf(
180+
SafetySetting(
181+
HarmCategory.HARASSMENT,
182+
HarmBlockThreshold.MEDIUM_AND_ABOVE,
183+
HarmBlockMethod.SEVERITY
184+
)
185+
)
186+
187+
val generativeModel =
188+
GenerativeModel(
189+
"gemini-2.5-flash",
190+
safetySettings = safetySettings,
191+
generativeBackend = GenerativeBackend.googleAI(),
192+
controller = apiController
193+
)
194+
195+
val exception =
196+
shouldThrow<InvalidStateException> {
197+
generativeModel.generateContent("my test prompt")
198+
}
199+
200+
exception.message shouldContain "HarmBlockMethod is unsupported by the Google Developer API"
201+
}
202+
203+
@Test
204+
fun `exception NOT thrown when using HarmBlockMethod with VertexAI`() = doBlocking {
205+
val mockEngine = MockEngine {
206+
respond(
207+
generateContentResponseAsJsonString("text response"),
208+
HttpStatusCode.OK,
209+
headersOf(HttpHeaders.ContentType, "application/json")
210+
)
211+
}
212+
213+
val apiController =
214+
APIController(
215+
"super_cool_test_key",
216+
"gemini-2.5-flash",
217+
RequestOptions(),
218+
mockEngine,
219+
TEST_CLIENT_ID,
220+
mockFirebaseApp,
221+
TEST_VERSION,
222+
TEST_APP_ID,
223+
null,
224+
)
225+
226+
val safetySettings =
227+
listOf(
228+
SafetySetting(
229+
HarmCategory.HARASSMENT,
230+
HarmBlockThreshold.MEDIUM_AND_ABOVE,
231+
HarmBlockMethod.SEVERITY
232+
)
233+
)
234+
235+
val generativeModel =
236+
GenerativeModel(
237+
"gemini-2.5-flash",
238+
safetySettings = safetySettings,
239+
generativeBackend = GenerativeBackend.vertexAI("us-central1"),
240+
controller = apiController
241+
)
242+
243+
withTimeout(5.seconds) { generativeModel.generateContent("my test prompt") }
244+
}
245+
149246
@OptIn(PublicPreviewAPI::class)
150247
private fun generateContentResponseAsJsonString(text: String): String {
151248
return JSON.encodeToString(

0 commit comments

Comments
 (0)