[Vertex AI] Swift Testing generateContentStream integration test (#14611)

Author: andrewheard

FirebaseVertexAI/Tests/TestApp/Sources/FirebaseAppUtils.swift

@@ -0,0 +1,37 @@
+// Copyright 2025 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+import FirebaseCore
+
+extension FirebaseApp {
+  /// Configures a Firebase app with the specified name and Google Service Info plist file name.
+  ///
+  /// - Parameters:
+  ///   - appName: The Firebase app's name; see ``FirebaseAppNames`` for app names with special
+  ///     meanings in the TestApp.
+  ///   - plistName: The file name of the Google Service Info plist, excluding the file extension;
+  ///     for the default app this is typically called `GoogleService-Info` but any file name may be
+  ///     used for other apps.
+  static func configure(appName: String, plistName: String) {
+    assert(!plistName.hasSuffix(".plist"), "The .plist file extension must be omitted.")
+    guard let plistPath =
+      Bundle.main.path(forResource: plistName, ofType: "plist") else {
+      fatalError("The file '\(plistName).plist' was not found.")
+    }
+    guard let options = FirebaseOptions(contentsOfFile: plistPath) else {
+      fatalError("Failed to parse options from '\(plistName).plist'.")
+    }
+    FirebaseApp.configure(name: appName, options: options)
+  }
+}

FirebaseVertexAI/Tests/TestApp/Sources/TestApp.swift

@@ -24,15 +24,15 @@ struct TestApp: App {
     // Configure default Firebase App
     FirebaseApp.configure()
 
+    // Configure a Firebase App that is the same as the default app but without App Check.
+    // This is used for tests that should fail when App Check is not configured.
+    FirebaseApp.configure(
+      appName: FirebaseAppNames.appCheckNotConfigured,
+      plistName: "GoogleService-Info"
+    )
+
     // Configure a Firebase App without a billing account (i.e., the "Spark" plan).
-    guard let plistPath =
-      Bundle.main.path(forResource: "GoogleService-Info-Spark", ofType: "plist") else {
-      fatalError("The file 'GoogleService-Info-Spark.plist' was not found.")
-    }
-    guard let options = FirebaseOptions(contentsOfFile: plistPath) else {
-      fatalError("Failed to parse options from 'GoogleService-Info-Spark.plist'.")
-    }
-    FirebaseApp.configure(name: FirebaseAppNames.spark, options: options)
+    FirebaseApp.configure(appName: FirebaseAppNames.spark, plistName: "GoogleService-Info-Spark")
   }
 
   var body: some Scene {

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

@@ -115,4 +115,81 @@ struct GenerateContentIntegrationTests {
     #expect(candidatesTokensDetails.modality == .text)
     #expect(candidatesTokensDetails.tokenCount == usageMetadata.candidatesTokenCount)
   }
+
+  // MARK: Streaming Tests
+
+  @Test(arguments: InstanceConfig.allConfigs)
+  func generateContentStream(_ config: InstanceConfig) async throws {
+    let expectedText = """
+    1. Mercury
+    2. Venus
+    3. Earth
+    4. Mars
+    5. Jupiter
+    6. Saturn
+    7. Uranus
+    8. Neptune
+    """
+    let prompt = """
+    What are the names of the planets in the solar system, ordered from closest to furthest from
+    the sun? Answer with a Markdown numbered list of the names and no other text.
+    """
+    let model = VertexAI.componentInstance(config).generativeModel(
+      modelName: ModelNames.gemini2FlashLite,
+      generationConfig: generationConfig,
+      safetySettings: safetySettings
+    )
+    let chat = model.startChat()
+
+    let stream = try chat.sendMessageStream(prompt)
+    var textValues = [String]()
+    for try await value in stream {
+      try textValues.append(#require(value.text))
+    }
+
+    let userHistory = try #require(chat.history.first)
+    #expect(userHistory.role == "user")
+    #expect(userHistory.parts.count == 1)
+    let promptTextPart = try #require(userHistory.parts.first as? TextPart)
+    #expect(promptTextPart.text == prompt)
+    let modelHistory = try #require(chat.history.last)
+    #expect(modelHistory.role == "model")
+    #expect(modelHistory.parts.count == 1)
+    let modelTextPart = try #require(modelHistory.parts.first as? TextPart)
+    let modelText = modelTextPart.text.trimmingCharacters(in: .whitespacesAndNewlines)
+    #expect(modelText == expectedText)
+    #expect(textValues.count > 1)
+    let text = textValues.joined().trimmingCharacters(in: .whitespacesAndNewlines)
+    #expect(text == expectedText)
+  }
+
+  // MARK: - App Check Tests
+
+  @Test(arguments: [
+    InstanceConfig.vertexV1AppCheckNotConfigured,
+    InstanceConfig.vertexV1BetaAppCheckNotConfigured,
+    // App Check is not supported on the Generative Language Developer API endpoint since it
+    // bypasses the Vertex AI in Firebase proxy.
+  ])
+  func generateContent_appCheckNotConfigured_shouldFail(_ config: InstanceConfig) async throws {
+    let model = VertexAI.componentInstance(config).generativeModel(
+      modelName: ModelNames.gemini2Flash
+    )
+    let prompt = "Where is Google headquarters located? Answer with the city name only."
+
+    try await #require {
+      _ = try await model.generateContent(prompt)
+    } throws: {
+      guard let error = $0 as? GenerateContentError else {
+        Issue.record("Expected a \(GenerateContentError.self); got \($0.self).")
+        return false
+      }
+      guard case let .internalError(underlyingError) = error else {
+        Issue.record("Expected a GenerateContentError.internalError(...); got \(error.self).")
+        return false
+      }
+
+      return String(describing: underlyingError).contains("Firebase App Check token is invalid")
+    }
+  }
 }

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

@@ -67,62 +67,6 @@ final class IntegrationTests: XCTestCase {
     storage = Storage.storage()
   }
 
-  // MARK: - Generate Content
-
-  func testGenerateContentStream() async throws {
-    let expectedText = """
-    1.  Mercury
-    2.  Venus
-    3.  Earth
-    4.  Mars
-    5.  Jupiter
-    6.  Saturn
-    7.  Uranus
-    8.  Neptune
-    """
-    let prompt = """
-    What are the names of the planets in the solar system, ordered from closest to furthest from
-    the sun? Answer with a Markdown numbered list of the names and no other text.
-    """
-    let chat = model.startChat()
-
-    let stream = try chat.sendMessageStream(prompt)
-    var textValues = [String]()
-    for try await value in stream {
-      try textValues.append(XCTUnwrap(value.text))
-    }
-
-    let userHistory = try XCTUnwrap(chat.history.first)
-    XCTAssertEqual(userHistory.role, "user")
-    XCTAssertEqual(userHistory.parts.count, 1)
-    let promptTextPart = try XCTUnwrap(userHistory.parts.first as? TextPart)
-    XCTAssertEqual(promptTextPart.text, prompt)
-    let modelHistory = try XCTUnwrap(chat.history.last)
-    XCTAssertEqual(modelHistory.role, "model")
-    XCTAssertEqual(modelHistory.parts.count, 1)
-    let modelTextPart = try XCTUnwrap(modelHistory.parts.first as? TextPart)
-    let modelText = modelTextPart.text.trimmingCharacters(in: .whitespacesAndNewlines)
-    XCTAssertEqual(modelText, expectedText)
-    XCTAssertGreaterThan(textValues.count, 1)
-    let text = textValues.joined().trimmingCharacters(in: .whitespacesAndNewlines)
-    XCTAssertEqual(text, expectedText)
-  }
-
-  func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
-    let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
-    addTeardownBlock { await app.delete() }
-    let vertex = VertexAI.vertexAI(app: app)
-    let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
-    let prompt = "Where is Google headquarters located? Answer with the city name only."
-
-    do {
-      _ = try await model.generateContent(prompt)
-      XCTFail("Expected a Firebase App Check error; none thrown.")
-    } catch let GenerateContentError.internalError(error) {
-      XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
-    }
-  }
-
   // MARK: - Count Tokens
 
   func testCountTokens_text() async throws {
@@ -285,8 +229,7 @@ final class IntegrationTests: XCTestCase {
   }
 
   func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
-    let app = try FirebaseApp.defaultNamedCopy(name: FirebaseAppNames.appCheckNotConfigured)
-    addTeardownBlock { await app.delete() }
+    let app = try XCTUnwrap(FirebaseApp.app(name: FirebaseAppNames.appCheckNotConfigured))
     let vertex = VertexAI.vertexAI(app: app)
     let model = vertex.generativeModel(modelName: "gemini-2.0-flash")
     let prompt = "Why is the sky blue?"

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

@@ -34,6 +34,15 @@ struct InstanceConfig {
   )
   static let allConfigs = [vertexV1, vertexV1Beta, developerV1, developerV1Beta]
 
+  static let vertexV1AppCheckNotConfigured = InstanceConfig(
+    appName: FirebaseAppNames.appCheckNotConfigured,
+    apiConfig: APIConfig(service: .vertexAI, version: .v1)
+  )
+  static let vertexV1BetaAppCheckNotConfigured = InstanceConfig(
+    appName: FirebaseAppNames.appCheckNotConfigured,
+    apiConfig: APIConfig(service: .vertexAI, version: .v1beta)
+  )
+
   let appName: String?
   let location: String?
   let apiConfig: APIConfig

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

@@ -21,8 +21,8 @@
 		8692F29A2CC9477800539E8F /* FirebaseAuth in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F2992CC9477800539E8F /* FirebaseAuth */; };
 		8692F29C2CC9477800539E8F /* FirebaseStorage in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29B2CC9477800539E8F /* FirebaseStorage */; };
 		8692F29E2CC9477800539E8F /* FirebaseVertexAI in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29D2CC9477800539E8F /* FirebaseVertexAI */; };
-		8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */; };
 		8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */; };
+		86CC31352D91EE9E0087E964 /* FirebaseAppUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */; };
 		86D77DFC2D7A5340003D155D /* GenerateContentIntegrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */; };
 		86D77DFE2D7B5C86003D155D /* GoogleService-Info-Spark.plist in Resources */ = {isa = PBXBuildFile; fileRef = 86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */; };
 		86D77E022D7B63AF003D155D /* Constants.swift in Sources */ = {isa = PBXBuildFile; fileRef = 86D77E012D7B63AC003D155D /* Constants.swift */; };
@@ -53,8 +53,8 @@
 		868A7C502CCC263300E449DD /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = "<group>"; };
 		868A7C532CCC26B500E449DD /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
 		868A7C552CCC271300E449DD /* TestApp.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TestApp.entitlements; sourceTree = "<group>"; };
-		8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppTestUtils.swift; sourceTree = "<group>"; };
 		8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestAppCheckProviderFactory.swift; sourceTree = "<group>"; };
+		86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppUtils.swift; sourceTree = "<group>"; };
 		86D77DFB2D7A5340003D155D /* GenerateContentIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = GenerateContentIntegrationTests.swift; sourceTree = "<group>"; };
 		86D77DFD2D7B5C86003D155D /* GoogleService-Info-Spark.plist */ = {isa = PBXFileReference; lastKnownFileType = text.plist.xml; path = "GoogleService-Info-Spark.plist"; sourceTree = "<group>"; };
 		86D77E012D7B63AC003D155D /* Constants.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = Constants.swift; sourceTree = "<group>"; };
@@ -129,6 +129,7 @@
 				8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */,
 				8661385D2CC943DD00F4B78E /* ContentView.swift */,
 				86D77E012D7B63AC003D155D /* Constants.swift */,
+				86CC31342D91EE9E0087E964 /* FirebaseAppUtils.swift */,
 			);
 			path = Sources;
 			sourceTree = "<group>";
@@ -158,7 +159,6 @@
 			isa = PBXGroup;
 			children = (
 				86D77E032D7B6C95003D155D /* InstanceConfig.swift */,
-				8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */,
 				862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */,
 			);
 			path = Utilities;
@@ -275,6 +275,7 @@
 			isa = PBXSourcesBuildPhase;
 			buildActionMask = 2147483647;
 			files = (
+				86CC31352D91EE9E0087E964 /* FirebaseAppUtils.swift in Sources */,
 				8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */,
 				8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */,
 				8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */,
@@ -288,7 +289,6 @@
 			files = (
 				8689CDCC2D7F8BD700BF426B /* CountTokensIntegrationTests.swift in Sources */,
 				86D77E042D7B6C9D003D155D /* InstanceConfig.swift in Sources */,
-				8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
 				868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
 				864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */,
 				862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */,