Skip to content

Commit 3c7c94c

Browse files
rlazorlazo
authored
[AI] Add integration tests for template models (#8001)
Introduced integration tests for `TemplateGenerativeModel` to validate its functionalities, with Chat features being tested in the `ChatTemplateIntegrationTests` file Updated `AIModels` to include `vertexAITemplateModel` and `googleAITemplateModel` for use in these tests.
1 parent 072e970 commit 3c7c94c

3 files changed

Lines changed: 211 additions & 0 deletions

File tree

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

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,9 @@ package com.google.firebase.ai
1818
import androidx.test.platform.app.InstrumentationRegistry
1919
import com.google.firebase.FirebaseApp
2020
import com.google.firebase.ai.type.GenerativeBackend
21+
import com.google.firebase.ai.type.PublicPreviewAPI
2122

23+
@OptIn(PublicPreviewAPI::class)
2224
class AIModels {
2325

2426
companion object {
@@ -31,6 +33,8 @@ class AIModels {
3133
lateinit var vertexAIFlashLiteModel: GenerativeModel
3234
lateinit var googleAIFlashModel: GenerativeModel
3335
lateinit var googleAIFlashLiteModel: GenerativeModel
36+
lateinit var vertexAITemplateModel: TemplateGenerativeModel
37+
lateinit var googleAITemplateModel: TemplateGenerativeModel
3438

3539
/** Returns a list of general purpose models to test */
3640
fun getModels(): List<GenerativeModel> {
@@ -45,6 +49,14 @@ class AIModels {
4549
)
4650
}
4751

52+
/** Returns a list of template models to test */
53+
fun getTemplateModels(): List<TemplateGenerativeModel> {
54+
if (app == null) {
55+
setup()
56+
}
57+
return listOf(vertexAITemplateModel, googleAITemplateModel)
58+
}
59+
4860
fun app(): FirebaseApp {
4961
if (app == null) {
5062
setup()
@@ -75,6 +87,10 @@ class AIModels {
7587
.generativeModel(
7688
modelName = "gemini-2.5-flash-lite",
7789
)
90+
vertexAITemplateModel =
91+
FirebaseAI.getInstance(app!!, GenerativeBackend.vertexAI()).templateGenerativeModel()
92+
googleAITemplateModel =
93+
FirebaseAI.getInstance(app!!, GenerativeBackend.googleAI()).templateGenerativeModel()
7894
}
7995
}
8096
}

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

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
/*
2+
* Copyright 2026 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+
17+
package com.google.firebase.ai
18+
19+
import com.google.firebase.ai.AIModels.Companion.getTemplateModels
20+
import com.google.firebase.ai.type.PublicPreviewAPI
21+
import com.google.firebase.ai.type.content
22+
import io.kotest.matchers.shouldBe
23+
import io.kotest.matchers.string.shouldContainIgnoringCase
24+
import kotlinx.coroutines.flow.toList
25+
import kotlinx.coroutines.runBlocking
26+
import org.junit.Test
27+
28+
@OptIn(PublicPreviewAPI::class)
29+
class ChatTemplateIntegrationTests {
30+
/*
31+
* Template used in these tests
32+
*
33+
* model: "<MODEL_NAME>"
34+
* input:
35+
* schema:
36+
* customerName: string, the name of the customer
37+
* topic: string, problem to solve
38+
* config:
39+
* temperature: 0.1
40+
* topK: 10
41+
* topP: 0.8
42+
*
43+
* -----
44+
* {{role "system"}}
45+
* You're a customer service agent, but you've been trained to only respond in very short
46+
* sentences. Be succinct, and to the point. No more than 5 words per response
47+
*
48+
* {{role "user"}}
49+
*
50+
* Hello, {{customerName}}
51+
*
52+
* Let's talk about {{topic}}
53+
* {{history}}
54+
*/
55+
private val templateId = "chat-test-template"
56+
57+
private val customerName = "John Doe"
58+
59+
private val topic = "Firebase"
60+
private val inputs = mapOf("customerName" to customerName, "topic" to topic)
61+
62+
@Test
63+
fun testTemplateChat_sendMessage() {
64+
for (model in getTemplateModels()) {
65+
runBlocking {
66+
val chat = model.startChat(templateId, inputs)
67+
val response = chat.sendMessage("which number is higher, one or ten?")
68+
69+
response.candidates.isNotEmpty() shouldBe true
70+
response.text shouldContainIgnoringCase "ten"
71+
72+
chat.history.size shouldBe 2
73+
}
74+
}
75+
}
76+
77+
@Test
78+
fun testTemplateChat_sendMessageStream() {
79+
for (model in getTemplateModels()) {
80+
runBlocking {
81+
val chat = model.startChat(templateId, inputs)
82+
val responses = chat.sendMessageStream("which number is higher, one or ten?").toList()
83+
responses.isNotEmpty() shouldBe true
84+
responses.joinToString { it.text ?: "" } shouldContainIgnoringCase "ten"
85+
chat.history.size shouldBe 2
86+
}
87+
}
88+
}
89+
90+
@Test
91+
fun testTemplateChat_withHistory() {
92+
for (model in getTemplateModels()) {
93+
runBlocking {
94+
val history =
95+
listOf(
96+
content("user") { text("which number is higher, one or ten?") },
97+
content("model") { text("Ten.") }
98+
)
99+
val chat = model.startChat(templateId, inputs, history)
100+
chat.history.size shouldBe 2
101+
val response =
102+
chat.sendMessage(
103+
"Please concatenate them both, first the smaller one, then the bigger one."
104+
)
105+
106+
response.candidates.isNotEmpty() shouldBe true
107+
response.text shouldContainIgnoringCase "oneten"
108+
109+
chat.history.size shouldBe 4
110+
}
111+
}
112+
}
113+
}

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

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
/*
2+
* Copyright 2026 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+
17+
package com.google.firebase.ai
18+
19+
import com.google.firebase.ai.AIModels.Companion.getTemplateModels
20+
import com.google.firebase.ai.type.PublicPreviewAPI
21+
import io.kotest.matchers.collections.shouldNotBeEmpty
22+
import io.kotest.matchers.string.shouldContainIgnoringCase
23+
import io.kotest.matchers.string.shouldNotBeEmpty
24+
import kotlinx.coroutines.flow.toList
25+
import kotlinx.coroutines.runBlocking
26+
import org.junit.Test
27+
28+
@OptIn(PublicPreviewAPI::class)
29+
class TemplateIntegrationTests {
30+
/*
31+
* Template used in these tests
32+
*
33+
* model: "<MODEL_NAME>"
34+
* input:
35+
* schema:
36+
* customerName: string, the name of the customer
37+
* topic: string, problem to solve
38+
* config:
39+
* temperature: 0.1
40+
* topK: 10
41+
* topP: 0.8
42+
*
43+
* -----
44+
*
45+
* Repeat back with "{{customerName}} - {{topic}}" and just that
46+
*/
47+
private val templateId = "test-template"
48+
49+
private val customerName = "John Doe"
50+
51+
private val topic = "Firebase"
52+
private val inputs = mapOf("customerName" to customerName, "topic" to topic)
53+
54+
@Test
55+
fun testTemplateGenerateContent() {
56+
for (model in getTemplateModels()) {
57+
runBlocking {
58+
val response = model.generateContent(templateId, inputs)
59+
60+
response.candidates.shouldNotBeEmpty()
61+
response.text shouldContainIgnoringCase customerName
62+
response.text shouldContainIgnoringCase topic
63+
}
64+
}
65+
}
66+
67+
@Test
68+
fun testTemplateGenerateContentStream() {
69+
for (model in getTemplateModels()) {
70+
runBlocking {
71+
val responses = model.generateContentStream(templateId, inputs).toList()
72+
responses
73+
.joinToString { it.text ?: "" }
74+
.lowercase()
75+
.let {
76+
it shouldContainIgnoringCase customerName
77+
it shouldContainIgnoringCase topic
78+
}
79+
}
80+
}
81+
}
82+
}

0 commit comments

Comments
 (0)