Skip to content

Commit f31257c

Browse files
Add tests for HarmBlockMethod compatibility with GoogleAI and VertexAI (#7567)
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. --------- Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>
1 parent e4369ab commit f31257c

File tree

1 file changed

+95
-0
lines changed

1 file changed

+95
-0
lines changed

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

Lines changed: 95 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,95 @@ 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> { generativeModel.generateContent("my test prompt") }
197+
198+
exception.message shouldContain "HarmBlockMethod is unsupported by the Google Developer API"
199+
}
200+
201+
@Test
202+
fun `exception NOT thrown when using HarmBlockMethod with VertexAI`() = doBlocking {
203+
val mockEngine = MockEngine {
204+
respond(
205+
generateContentResponseAsJsonString("text response"),
206+
HttpStatusCode.OK,
207+
headersOf(HttpHeaders.ContentType, "application/json")
208+
)
209+
}
210+
211+
val apiController =
212+
APIController(
213+
"super_cool_test_key",
214+
"gemini-2.5-flash",
215+
RequestOptions(),
216+
mockEngine,
217+
TEST_CLIENT_ID,
218+
mockFirebaseApp,
219+
TEST_VERSION,
220+
TEST_APP_ID,
221+
null,
222+
)
223+
224+
val safetySettings =
225+
listOf(
226+
SafetySetting(
227+
HarmCategory.HARASSMENT,
228+
HarmBlockThreshold.MEDIUM_AND_ABOVE,
229+
HarmBlockMethod.SEVERITY
230+
)
231+
)
232+
233+
val generativeModel =
234+
GenerativeModel(
235+
"gemini-2.5-flash",
236+
safetySettings = safetySettings,
237+
generativeBackend = GenerativeBackend.vertexAI("us-central1"),
238+
controller = apiController
239+
)
240+
241+
withTimeout(5.seconds) { generativeModel.generateContent("my test prompt") }
242+
}
243+
149244
@OptIn(PublicPreviewAPI::class)
150245
private fun generateContentResponseAsJsonString(text: String): String {
151246
return JSON.encodeToString(

0 commit comments

Comments
 (0)