Skip to content

Commit 6191e9d

Browse files
andrewheardandrewheard
authored
[Vertex AI] Swift Testing generateContentStream integration test (#14611)
1 parent 70c8908 commit 6191e9d

File tree

7 files changed

+136
-110
lines changed

7 files changed

+136
-110
lines changed

‎FirebaseVertexAI/Tests/TestApp/Sources/FirebaseAppUtils.swift

+37
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
// Copyright 2025 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import FirebaseCore
16+
17+
extension FirebaseApp {
18+
/// Configures a Firebase app with the specified name and Google Service Info plist file name.
19+
///
20+
/// - Parameters:
21+
/// - appName: The Firebase app's name; see ``FirebaseAppNames`` for app names with special
22+
/// meanings in the TestApp.
23+
/// - plistName: The file name of the Google Service Info plist, excluding the file extension;
24+
/// for the default app this is typically called `GoogleService-Info` but any file name may be
25+
/// used for other apps.
26+
static func configure(appName: String, plistName: String) {
27+
assert(!plistName.hasSuffix(".plist"), "The .plist file extension must be omitted.")
28+
guard let plistPath =
29+
Bundle.main.path(forResource: plistName, ofType: "plist") else {
30+
fatalError("The file '\(plistName).plist' was not found.")
31+
}
32+
guard let options = FirebaseOptions(contentsOfFile: plistPath) else {
33+
fatalError("Failed to parse options from '\(plistName).plist'.")
34+
}
35+
FirebaseApp.configure(name: appName, options: options)
36+
}
37+
}

‎FirebaseVertexAI/Tests/TestApp/Sources/TestApp.swift

+8-8
Original file line numberDiff line numberDiff line change
@@ -24,15 +24,15 @@ struct TestApp: App {
2424
// Configure default Firebase App
2525
FirebaseApp.configure()
2626

27+
// Configure a Firebase App that is the same as the default app but without App Check.
28+
// This is used for tests that should fail when App Check is not configured.
29+
FirebaseApp.configure(
30+
appName: FirebaseAppNames.appCheckNotConfigured,
31+
plistName: "GoogleService-Info"
32+
)
33+
2734
// Configure a Firebase App without a billing account (i.e., the "Spark" plan).
28-
guard let plistPath =
29-
Bundle.main.path(forResource: "GoogleService-Info-Spark", ofType: "plist") else {
30-
fatalError("The file 'GoogleService-Info-Spark.plist' was not found.")
31-
}
32-
guard let options = FirebaseOptions(contentsOfFile: plistPath) else {
33-
fatalError("Failed to parse options from 'GoogleService-Info-Spark.plist'.")
34-
}
35-
FirebaseApp.configure(name: FirebaseAppNames.spark, options: options)
35+
FirebaseApp.configure(appName: FirebaseAppNames.spark, plistName: "GoogleService-Info-Spark")
3636
}
3737

3838
var body: some Scene {

‎FirebaseVertexAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift

+77
Original file line numberDiff line numberDiff line change
@@ -115,4 +115,81 @@ struct GenerateContentIntegrationTests {
115115
#expect(candidatesTokensDetails.modality == .text)
116116
#expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
117117
}
118+
119+
// MARK: Streaming Tests
120+
121+
@Test(arguments: InstanceConfig.allConfigs)
122+
func generateContentStream(_ config: InstanceConfig) async throws {
123+
let expectedText = """
124+
1. Mercury
125+
2. Venus
126+
3. Earth
127+
4. Mars
128+
5. Jupiter
129+
6. Saturn
130+
7. Uranus
131+
8. Neptune
132+
"""
133+
let prompt = """
134+
What are the names of the planets in the solar system, ordered from closest to furthest from
135+
the sun? Answer with a Markdown numbered list of the names and no other text.
136+
"""
137+
let model = VertexAI.componentInstance(config).generativeModel(
138+
modelName: ModelNames.gemini2FlashLite,
139+
generationConfig: generationConfig,
140+
safetySettings: safetySettings
141+
)
142+
let chat = model.startChat()
143+
144+
let stream = try chat.sendMessageStream(prompt)
145+
var textValues = [String]()
146+
for try await value in stream {
147+
try textValues.append(#require(value.text))
148+
}
149+
150+
let userHistory = try #require(chat.history.first)
151+
#expect(userHistory.role == "user")
152+
#expect(userHistory.parts.count == 1)
153+
let promptTextPart = try #require(userHistory.parts.first as? TextPart)
154+
#expect(promptTextPart.text == prompt)
155+
let modelHistory = try #require(chat.history.last)
156+
#expect(modelHistory.role == "model")
157+
#expect(modelHistory.parts.count == 1)
158+
let modelTextPart = try #require(modelHistory.parts.first as? TextPart)
159+
let modelText = modelTextPart.text.trimmingCharacters(in: .whitespacesAndNewlines)
160+
#expect(modelText == expectedText)
161+
#expect(textValues.count > 1)
162+
let text = textValues.joined().trimmingCharacters(in: .whitespacesAndNewlines)
163+
#expect(text == expectedText)
164+
}
165+
166+
// MARK: - App Check Tests
167+
168+
@Test(arguments: [
169+
InstanceConfig.vertexV1AppCheckNotConfigured,
170+
InstanceConfig.vertexV1BetaAppCheckNotConfigured,
171+
// App Check is not supported on the Generative Language Developer API endpoint since it
172+
// bypasses the Vertex AI in Firebase proxy.
173+
])
174+
func generateContent_appCheckNotConfigured_shouldFail(_ config: InstanceConfig) async throws {
175+
let model = VertexAI.componentInstance(config).generativeModel(
176+
modelName: ModelNames.gemini2Flash
177+
)
178+
let prompt = "Where is Google headquarters located? Answer with the city name only."
179+
180+
try await #require {
181+
_ = try await model.generateContent(prompt)
182+
} throws: {
183+
guard let error = $0 as? GenerateContentError else {
184+
Issue.record("Expected a \(GenerateContentError.self); got \($0.self).")
185+
return false
186+
}
187+
guard case let .internalError(underlyingError) = error else {
188+
Issue.record("Expected a GenerateContentError.internalError(...); got \(error.self).")
189+
return false
190+
}
191+
192+
return String(describing: underlyingError).contains("Firebase App Check token is invalid")
193+
}
194+
}
118195
}

‎FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

+1-58
Original file line numberDiff line numberDiff line change
@@ -67,62 +67,6 @@ final class IntegrationTests: XCTestCase {
6767
storage = Storage.storage()
6868
}
6969

70-
// MARK: - Generate Content
71-
72-
func testGenerateContentStream() async throws {
73-
let expectedText = """
74-
1. Mercury
75-
2. Venus
76-
3. Earth
77-
4. Mars
78-
5. Jupiter
79-
6. Saturn
80-
7. Uranus
81-
8. Neptune
82-
"""
83-
let prompt = """
84-
What are the names of the planets in the solar system, ordered from closest to furthest from
85-
the sun? Answer with a Markdown numbered list of the names and no other text.
86-
"""
87-
let chat = model.startChat()
88-
89-
let stream = try chat.sendMessageStream(prompt)
90-
var textValues = [String]()
91-
for try await value in stream {
92-
try textValues.append(XCTUnwrap(value.text))
93-
}
94-
95-
let userHistory = try XCTUnwrap(chat.history.first)
96-
XCTAssertEqual(userHistory.role, "user")
97-
XCTAssertEqual(userHistory.parts.count, 1)
98-
let promptTextPart = try XCTUnwrap(userHistory.parts.first as? TextPart)
99-
XCTAssertEqual(promptTextPart.text, prompt)
100-
let modelHistory = try XCTUnwrap(chat.history.last)
101-
XCTAssertEqual(modelHistory.role, "model")
102-
XCTAssertEqual(modelHistory.parts.count, 1)
103-
let modelTextPart = try XCTUnwrap(modelHistory.parts.first as? TextPart)
104-
let modelText = modelTextPart.text.trimmingCharacters(in: .whitespacesAndNewlines)
105-
XCTAssertEqual(modelText, expectedText)
106-
XCTAssertGreaterThan(textValues.count, 1)
107-
let text = textValues.joined().trimmingCharacters(in: .whitespacesAndNewlines)
108-
XCTAssertEqual(text, expectedText)
109-
}
110-
111-
func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
112-
let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
113-
addTeardownBlock { await app.delete() }
114-
let vertex = VertexAI.vertexAI(app: app)
115-
let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
116-
let prompt = "Where is Google headquarters located? Answer with the city name only."
117-
118-
do {
119-
_ = try await model.generateContent(prompt)
120-
XCTFail("Expected a Firebase App Check error; none thrown.")
121-
} catch let GenerateContentError.internalError(error) {
122-
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
123-
}
124-
}
125-
12670
// MARK: - Count Tokens
12771

12872
func testCountTokens_text() async throws {
@@ -285,8 +229,7 @@ final class IntegrationTests: XCTestCase {
285229
}
286230

287231
func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
288-
let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
289-
addTeardownBlock { await app.delete() }
232+
let app = try XCTUnwrap(FirebaseApp.app(name: FirebaseAppNames.appCheckNotConfigured))
290233
let vertex = VertexAI.vertexAI(app: app)
291234
let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
292235
let prompt = "Why is the sky blue?"

‎FirebaseVertexAI/Tests/TestApp/Tests/Utilities/FirebaseAppTestUtils.swift

-40
This file was deleted.

‎FirebaseVertexAI/Tests/TestApp/Tests/Utilities/InstanceConfig.swift

+9
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,15 @@ struct InstanceConfig {
3434
)
3535
static let allConfigs = [vertexV1, vertexV1Beta, developerV1, developerV1Beta]
3636

37+
static let vertexV1AppCheckNotConfigured = InstanceConfig(
38+
appName: FirebaseAppNames.appCheckNotConfigured,
39+
apiConfig: APIConfig(service: .vertexAI, version: .v1)
40+
)
41+
static let vertexV1BetaAppCheckNotConfigured = InstanceConfig(
42+
appName: FirebaseAppNames.appCheckNotConfigured,
43+
apiConfig: APIConfig(service: .vertexAI, version: .v1beta)
44+
)
45+
3746
let appName: String?
3847
let location: String?
3948
let apiConfig: APIConfig

‎FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj

+4-4
Original file line numberDiff line numberDiff line change
@@ -21,8 +21,8 @@
2121
8692F29A2CC9477800539E8F /* FirebaseAuth in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F2992CC9477800539E8F /* FirebaseAuth */; };
2222
8692F29C2CC9477800539E8F /* FirebaseStorage in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29B2CC9477800539E8F /* FirebaseStorage */; };
2323
8692F29E2CC9477800539E8F /* FirebaseVertexAI in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29D2CC9477800539E8F /* FirebaseVertexAI */; };
24-
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */; };
2524
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */; };
25+
86CC31352D91EE9E0087E964 /* FirebaseAppUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */; };
2626
86D77DFC2D7A5340003D155D /* GenerateContentIntegrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */; };
2727
86D77DFE2D7B5C86003D155D /* GoogleService-Info-Spark.plist in Resources */ = {isa = PBXBuildFile; fileRef = 86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */; };
2828
86D77E022D7B63AF003D155D /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E012D7B63AC003D155D /* Constants.swift */; };
@@ -53,8 +53,8 @@
5353
868A7C502CCC263300E449DD /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = "<group>"; };
5454
868A7C532CCC26B500E449DD /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
5555
868A7C552CCC271300E449DD /* TestApp.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TestApp.entitlements; sourceTree = "<group>"; };
56-
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppTestUtils.swift; sourceTree = "<group>"; };
5756
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestAppCheckProviderFactory.swift; sourceTree = "<group>"; };
57+
86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppUtils.swift; sourceTree = "<group>"; };
5858
86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = GenerateContentIntegrationTests.swift; sourceTree = "<group>"; };
5959
86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = "GoogleService-Info-Spark.plist"; sourceTree = "<group>"; };
6060
86D77E012D7B63AC003D155D /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = "<group>"; };
@@ -129,6 +129,7 @@
129129
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */,
130130
8661385D2CC943DD00F4B78E /* ContentView.swift */,
131131
86D77E012D7B63AC003D155D /* Constants.swift */,
132+
86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */,
132133
);
133134
path = Sources;
134135
sourceTree = "<group>";
@@ -158,7 +159,6 @@
158159
isa = PBXGroup;
159160
children = (
160161
86D77E032D7B6C95003D155D /* InstanceConfig.swift */,
161-
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */,
162162
862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */,
163163
);
164164
path = Utilities;
@@ -275,6 +275,7 @@
275275
isa = PBXSourcesBuildPhase;
276276
buildActionMask = 2147483647;
277277
files = (
278+
86CC31352D91EE9E0087E964 /* FirebaseAppUtils.swift in Sources */,
278279
8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */,
279280
8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */,
280281
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */,
@@ -288,7 +289,6 @@
288289
files = (
289290
8689CDCC2D7F8BD700BF426B /* CountTokensIntegrationTests.swift in Sources */,
290291
86D77E042D7B6C9D003D155D /* InstanceConfig.swift in Sources */,
291-
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
292292
868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
293293
864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */,
294294
862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */,

0 commit comments

Comments
 (0)