Skip to content

Commit ece0efd

Browse files
emilypgoogleemilypgoogle
authored
Add AI integration tests (#7038)
Currently lacking mechanism to get API key from GitHub workflow and the workflow to regularly run the tests. This functionality will be added in a follow up PR. Includes a modification to an internal serializable type to better interact with failure cases and prevent crashes.
1 parent 3b7bfef commit ece0efd

File tree

7 files changed

+723
-3
lines changed

7 files changed

+723
-3
lines changed

‎firebase-ai/firebase-ai.gradle.kts

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -67,7 +67,10 @@ android {
6767
targetSdk = targetSdkVersion
6868
baseline = file("lint-baseline.xml")
6969
}
70-
sourceSets { getByName("test").java.srcDirs("src/testUtil") }
70+
sourceSets {
71+
// getByName("test").java.srcDirs("src/testUtil")
72+
getByName("androidTest") { kotlin.srcDirs("src/testUtil") }
73+
}
7174
}
7275

7376
// Enable Kotlin "Explicit API Mode". This causes the Kotlin compiler to fail if any

‎firebase-ai/src/androidTest/kotlin/com/google/firebase/ai/AIModels.kt

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.firebase.ai
17+
18+
import androidx.test.platform.app.InstrumentationRegistry
19+
import com.google.firebase.FirebaseApp
20+
import com.google.firebase.FirebaseOptions
21+
import com.google.firebase.ai.type.GenerativeBackend
22+
23+
class AIModels {
24+
25+
companion object {
26+
private val API_KEY: String = ""
27+
private val APP_ID: String = ""
28+
private val PROJECT_ID: String = "fireescape-integ-tests"
29+
// General purpose models
30+
var app: FirebaseApp? = null
31+
var flash2Model: GenerativeModel? = null
32+
var flash2LiteModel: GenerativeModel? = null
33+
34+
/** Returns a list of general purpose models to test */
35+
fun getModels(): List<GenerativeModel> {
36+
if (flash2Model == null) {
37+
setup()
38+
}
39+
return listOf(flash2Model!!, flash2LiteModel!!)
40+
}
41+
42+
fun app(): FirebaseApp {
43+
if (app == null) {
44+
setup()
45+
}
46+
return app!!
47+
}
48+
49+
fun setup() {
50+
val context = InstrumentationRegistry.getInstrumentation().context
51+
app =
52+
FirebaseApp.initializeApp(
53+
context,
54+
FirebaseOptions.Builder()
55+
.setApiKey(API_KEY)
56+
.setApplicationId(APP_ID)
57+
.setProjectId(PROJECT_ID)
58+
.build()
59+
)
60+
flash2Model =
61+
FirebaseAI.getInstance(app!!, GenerativeBackend.vertexAI())
62+
.generativeModel(
63+
modelName = "gemini-2.0-flash",
64+
)
65+
flash2LiteModel =
66+
FirebaseAI.getInstance(app!!, GenerativeBackend.vertexAI())
67+
.generativeModel(
68+
modelName = "gemini-2.0-flash-lite",
69+
)
70+
}
71+
}
72+
}

‎firebase-ai/src/androidTest/kotlin/com/google/firebase/ai/CountTokensTests.kt

Lines changed: 204 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,204 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.firebase.ai
17+
18+
import android.graphics.Bitmap
19+
import com.google.firebase.ai.AIModels.Companion.getModels
20+
import com.google.firebase.ai.type.Content
21+
import com.google.firebase.ai.type.ContentModality
22+
import com.google.firebase.ai.type.CountTokensResponse
23+
import java.io.ByteArrayOutputStream
24+
import kotlinx.coroutines.runBlocking
25+
import org.junit.Test
26+
27+
class CountTokensTests {
28+
29+
/** Ensures that the token count is expected for simple words. */
30+
@Test
31+
fun testCountTokensAmount() {
32+
for (model in getModels()) {
33+
runBlocking {
34+
val response = model.countTokens("this is five different words")
35+
assert(response.totalTokens == 5)
36+
assert(response.promptTokensDetails.size == 1)
37+
assert(response.promptTokensDetails[0].modality == ContentModality.TEXT)
38+
assert(response.promptTokensDetails[0].tokenCount == 5)
39+
}
40+
}
41+
}
42+
43+
/** Ensures that the model returns token counts in the correct modality for text. */
44+
@Test
45+
fun testCountTokensTextModality() {
46+
for (model in getModels()) {
47+
runBlocking {
48+
val response = model.countTokens("this is a text prompt")
49+
checkTokenCountsMatch(response)
50+
assert(response.promptTokensDetails.size == 1)
51+
assert(containsModality(response, ContentModality.TEXT))
52+
}
53+
}
54+
}
55+
56+
/** Ensures that the model returns token counts in the correct modality for bitmap images. */
57+
@Test
58+
fun testCountTokensImageModality() {
59+
for (model in getModels()) {
60+
runBlocking {
61+
val bitmap = Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888)
62+
val response = model.countTokens(bitmap)
63+
checkTokenCountsMatch(response)
64+
assert(response.promptTokensDetails.size == 1)
65+
assert(containsModality(response, ContentModality.IMAGE))
66+
}
67+
}
68+
}
69+
70+
/**
71+
* Ensures the model can count tokens for multiple modalities at once, and return the
72+
* corresponding token modalities correctly.
73+
*/
74+
@Test
75+
fun testCountTokensTextAndImageModality() {
76+
for (model in getModels()) {
77+
runBlocking {
78+
val bitmap = Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888)
79+
val response =
80+
model.countTokens(
81+
Content.Builder().text("this is text").build(),
82+
Content.Builder().image(bitmap).build()
83+
)
84+
checkTokenCountsMatch(response)
85+
assert(response.promptTokensDetails.size == 2)
86+
assert(containsModality(response, ContentModality.TEXT))
87+
assert(containsModality(response, ContentModality.IMAGE))
88+
}
89+
}
90+
}
91+
92+
/**
93+
* Ensures the model can count the tokens for a sent file. Additionally, ensures that the model
94+
* treats this sent file as the modality of the mime type, in this case, a plaintext file has its
95+
* tokens counted as `ContentModality.TEXT`.
96+
*/
97+
@Test
98+
fun testCountTokensTextFileModality() {
99+
for (model in getModels()) {
100+
runBlocking {
101+
val response =
102+
model.countTokens(
103+
Content.Builder().inlineData("this is text".toByteArray(), "text/plain").build()
104+
)
105+
checkTokenCountsMatch(response)
106+
assert(response.totalTokens == 3)
107+
assert(response.promptTokensDetails.size == 1)
108+
assert(containsModality(response, ContentModality.TEXT))
109+
}
110+
}
111+
}
112+
113+
/**
114+
* Ensures the model can count the tokens for a sent file. Additionally, ensures that the model
115+
* treats this sent file as the modality of the mime type, in this case, a PNG encoded bitmap has
116+
* its tokens counted as `ContentModality.IMAGE`.
117+
*/
118+
@Test
119+
fun testCountTokensImageFileModality() {
120+
for (model in getModels()) {
121+
runBlocking {
122+
val bitmap = Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888)
123+
val stream = ByteArrayOutputStream()
124+
bitmap.compress(Bitmap.CompressFormat.PNG, 1, stream)
125+
val array = stream.toByteArray()
126+
val response = model.countTokens(Content.Builder().inlineData(array, "image/png").build())
127+
checkTokenCountsMatch(response)
128+
assert(response.promptTokensDetails.size == 1)
129+
assert(containsModality(response, ContentModality.IMAGE))
130+
}
131+
}
132+
}
133+
134+
/**
135+
* Ensures that nothing is free, that is, empty content contains no tokens. For some reason, this
136+
* is treated as `ContentModality.TEXT`.
137+
*/
138+
@Test
139+
fun testCountTokensNothingIsFree() {
140+
for (model in getModels()) {
141+
runBlocking {
142+
val response = model.countTokens(Content.Builder().build())
143+
checkTokenCountsMatch(response)
144+
assert(response.totalTokens == 0)
145+
assert(response.promptTokensDetails.size == 1)
146+
assert(containsModality(response, ContentModality.TEXT))
147+
}
148+
}
149+
}
150+
151+
/**
152+
* Checks if the model can count the tokens for a sent file. Additionally, ensures that the model
153+
* treats this sent file as the modality of the mime type, in this case, a JSON file is not
154+
* recognized, and no tokens are counted. This ensures if/when the model can handle JSON, our
155+
* testing makes us aware.
156+
*/
157+
@Test
158+
fun testCountTokensJsonFileModality() {
159+
for (model in getModels()) {
160+
runBlocking {
161+
val json =
162+
"""
163+
{
164+
"foo": "bar",
165+
"baz": 3,
166+
"qux": [
167+
{
168+
"quux": [
169+
1,
170+
2
171+
]
172+
}
173+
]
174+
}
175+
"""
176+
.trimIndent()
177+
val response =
178+
model.countTokens(
179+
Content.Builder().inlineData(json.toByteArray(), "application/json").build()
180+
)
181+
checkTokenCountsMatch(response)
182+
assert(response.promptTokensDetails.isEmpty())
183+
assert(response.totalTokens == 0)
184+
}
185+
}
186+
}
187+
188+
fun checkTokenCountsMatch(response: CountTokensResponse) {
189+
assert(sumTokenCount(response) == response.totalTokens)
190+
}
191+
192+
fun sumTokenCount(response: CountTokensResponse): Int {
193+
return response.promptTokensDetails.sumOf { it.tokenCount }
194+
}
195+
196+
fun containsModality(response: CountTokensResponse, modality: ContentModality): Boolean {
197+
for (token in response.promptTokensDetails) {
198+
if (token.modality == modality) {
199+
return true
200+
}
201+
}
202+
return false
203+
}
204+
}

‎firebase-ai/src/androidTest/kotlin/com/google/firebase/ai/GenerateContentTests.kt

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,83 @@
1+
/*
2+
* Copyright 2025 Google LLC
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License");
5+
* you may not use this file except in compliance with the License.
6+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
package com.google.firebase.ai
17+
18+
import android.graphics.Bitmap
19+
import com.google.firebase.ai.AIModels.Companion.getModels
20+
import com.google.firebase.ai.type.Content
21+
import kotlinx.coroutines.runBlocking
22+
import org.junit.Test
23+
24+
class GenerateContentTests {
25+
private val validator = TypesValidator()
26+
27+
/**
28+
* Ensures the model can response to prompts and that the structure of this response is expected.
29+
*/
30+
@Test
31+
fun testGenerateContent_BasicRequest() {
32+
for (model in getModels()) {
33+
runBlocking {
34+
val response = model.generateContent("pick a random color")
35+
validator.validateResponse(response)
36+
}
37+
}
38+
}
39+
40+
/**
41+
* Ensures that the model can answer very simple questions. Further testing the "logic" of the
42+
* model and the content of the responses is prone to flaking, this test is also prone to that.
43+
* This is probably the furthest we can consistently test for reasonable response structure, past
44+
* sending the request and response back to the model and asking it if it fits our expectations.
45+
*/
46+
@Test
47+
fun testGenerateContent_ColorMixing() {
48+
for (model in getModels()) {
49+
runBlocking {
50+
val response = model.generateContent("what color is created when red and yellow are mixed?")
51+
validator.validateResponse(response)
52+
assert(response.text!!.contains("orange", true))
53+
}
54+
}
55+
}
56+
57+
/**
58+
* Ensures that the model can answer very simple questions. Further testing the "logic" of the
59+
* model and the content of the responses is prone to flaking, this test is also prone to that.
60+
* This is probably the furthest we can consistently test for reasonable response structure, past
61+
* sending the request and response back to the model and asking it if it fits our expectations.
62+
*/
63+
@Test
64+
fun testGenerateContent_CanSendImage() {
65+
for (model in getModels()) {
66+
runBlocking {
67+
val bitmap = Bitmap.createBitmap(10, 10, Bitmap.Config.ARGB_8888)
68+
val yellow = Integer.parseUnsignedInt("FFFFFF00", 16)
69+
bitmap.setPixel(3, 3, yellow)
70+
bitmap.setPixel(6, 3, yellow)
71+
bitmap.setPixel(3, 6, yellow)
72+
bitmap.setPixel(4, 7, yellow)
73+
bitmap.setPixel(5, 7, yellow)
74+
bitmap.setPixel(6, 6, yellow)
75+
val response =
76+
model.generateContent(
77+
Content.Builder().text("here is a tiny smile").image(bitmap).build()
78+
)
79+
validator.validateResponse(response)
80+
}
81+
}
82+
}
83+
}

0 commit comments

Comments
 (0)