Skip to content

Commit 5205d0d

Browse files
andrewheardandrewheard
committed
[Vertex AI] Add integration tests for App Check failure
1 parent 33a44a9 commit 5205d0d

File tree

5 files changed

+112
-1
lines changed

5 files changed

+112
-1
lines changed

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,7 @@ import SwiftUI
1919
@main
2020
struct TestApp: App {
2121
init() {
22-
AppCheck.setAppCheckProviderFactory(AppCheckDebugProviderFactory())
22+
AppCheck.setAppCheckProviderFactory(TestAppCheckProviderFactory())
2323
FirebaseApp.configure()
2424
}
2525

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

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,22 @@
1+
import FirebaseAppCheck
2+
import FirebaseCore
3+
import Foundation
4+
5+
/// An `AppCheckProviderFactory` for the Test App.
6+
///
7+
/// Defaults to the `AppCheckDebugProvider` unless the `FirebaseApp` `name` contains
8+
/// ``notConfiguredName``, in which case App Check is not configured; this facilitates integration
9+
/// testing of App Check failure cases.
10+
public class TestAppCheckProviderFactory: NSObject, AppCheckProviderFactory {
11+
/// The name, or a substring of the name, of Firebase apps where App Check is not configured.
12+
public static let notConfiguredName = "app-check-not-configured"
13+
14+
/// Returns the `AppCheckDebugProvider` unless `app.name` contains ``notConfiguredName``.
15+
public func createProvider(with app: FirebaseApp) -> (any AppCheckProvider)? {
16+
if app.name.contains(TestAppCheckProviderFactory.notConfiguredName) {
17+
return nil
18+
}
19+
20+
return AppCheckDebugProvider(app: app)
21+
}
22+
}

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

Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@ import FirebaseAuth
1616
import FirebaseCore
1717
import FirebaseStorage
1818
import FirebaseVertexAI
19+
import VertexAITestApp
1920
import XCTest
2021

2122
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
@@ -75,6 +76,21 @@ final class IntegrationTests: XCTestCase {
7576
XCTAssertEqual(text, "Mountain View")
7677
}
7778

79+
func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
80+
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
81+
addTeardownBlock { await app.delete() }
82+
let vertex = VertexAI.vertexAI(app: app)
83+
let model = vertex.generativeModel(modelName: "gemini-1.5-flash")
84+
let prompt = "Where is Google headquarters located? Answer with the city name only."
85+
86+
do {
87+
_ = try await model.generateContent(prompt)
88+
XCTFail("Expected a Firebase App Check error; none thrown.")
89+
} catch let GenerateContentError.internalError(error) {
90+
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
91+
}
92+
}
93+
7894
// MARK: - Count Tokens
7995

8096
func testCountTokens_text() async throws {
@@ -205,6 +221,21 @@ final class IntegrationTests: XCTestCase {
205221
XCTAssertEqual(response.totalTokens, 34)
206222
XCTAssertEqual(response.totalBillableCharacters, 59)
207223
}
224+
225+
func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
226+
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
227+
addTeardownBlock { await app.delete() }
228+
let vertex = VertexAI.vertexAI(app: app)
229+
let model = vertex.generativeModel(modelName: "gemini-1.5-flash")
230+
let prompt = "Why is the sky blue?"
231+
232+
do {
233+
_ = try await model.countTokens(prompt)
234+
XCTFail("Expected a Firebase App Check error; none thrown.")
235+
} catch {
236+
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
237+
}
238+
}
208239
}
209240

210241
extension StorageReference {

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

Lines changed: 40 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Copyright 2024 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 another `FirebaseApp` with the specified `name` and the same `FirebaseOptions`.
19+
func namedCopy(name: String) throws -> FirebaseApp {
20+
FirebaseApp.configure(name: name, options: options)
21+
guard let app = FirebaseApp.app(name: name) else {
22+
throw AppNotFound(name: name)
23+
}
24+
return app
25+
}
26+
27+
/// Configures an app with the specified `name` and the same `FirebaseOptions` as the default app.
28+
static func defaultNamedCopy(name: String) throws -> FirebaseApp {
29+
guard FirebaseApp.isDefaultAppConfigured(), let defaultApp = FirebaseApp.app() else {
30+
throw DefaultAppNotConfigured()
31+
}
32+
return try defaultApp.namedCopy(name: name)
33+
}
34+
35+
struct AppNotFound: Error {
36+
let name: String
37+
}
38+
39+
struct DefaultAppNotConfigured: Error {}
40+
}

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

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,8 @@
1818
8692F29A2CC9477800539E8F /* FirebaseAuth in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F2992CC9477800539E8F /* FirebaseAuth */; };
1919
8692F29C2CC9477800539E8F /* FirebaseStorage in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29B2CC9477800539E8F /* FirebaseStorage */; };
2020
8692F29E2CC9477800539E8F /* FirebaseVertexAI in Frameworks */ = {isa = PBXBuildFile; productRef = 8692F29D2CC9477800539E8F /* FirebaseVertexAI */; };
21+
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */; };
22+
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */; };
2123
/* End PBXBuildFile section */
2224

2325
/* Begin PBXContainerItemProxy section */
@@ -41,6 +43,8 @@
4143
868A7C502CCC263300E449DD /* Preview Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = "Preview Assets.xcassets"; sourceTree = "<group>"; };
4244
868A7C532CCC26B500E449DD /* Assets.xcassets */ = {isa = PBXFileReference; lastKnownFileType = folder.assetcatalog; path = Assets.xcassets; sourceTree = "<group>"; };
4345
868A7C552CCC271300E449DD /* TestApp.entitlements */ = {isa = PBXFileReference; lastKnownFileType = text.plist.entitlements; path = TestApp.entitlements; sourceTree = "<group>"; };
46+
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = FirebaseAppTestUtils.swift; sourceTree = "<group>"; };
47+
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestAppCheckProviderFactory.swift; sourceTree = "<group>"; };
4448
/* End PBXFileReference section */
4549

4650
/* Begin PBXFrameworksBuildPhase section */
@@ -107,6 +111,7 @@
107111
isa = PBXGroup;
108112
children = (
109113
8661385B2CC943DD00F4B78E /* TestApp.swift */,
114+
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */,
110115
8661385D2CC943DD00F4B78E /* ContentView.swift */,
111116
);
112117
path = Sources;
@@ -125,10 +130,19 @@
125130
isa = PBXGroup;
126131
children = (
127132
868A7C572CCC27AF00E449DD /* Integration */,
133+
8698D7442CD3CEF700ABA833 /* Utilities */,
128134
);
129135
path = Tests;
130136
sourceTree = "<group>";
131137
};
138+
8698D7442CD3CEF700ABA833 /* Utilities */ = {
139+
isa = PBXGroup;
140+
children = (
141+
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */,
142+
);
143+
path = Utilities;
144+
sourceTree = "<group>";
145+
};
132146
/* End PBXGroup section */
133147

134148
/* Begin PBXNativeTarget section */
@@ -241,13 +255,15 @@
241255
files = (
242256
8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */,
243257
8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */,
258+
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */,
244259
);
245260
runOnlyForDeploymentPostprocessing = 0;
246261
};
247262
866138652CC943DE00F4B78E /* Sources */ = {
248263
isa = PBXSourcesBuildPhase;
249264
buildActionMask = 2147483647;
250265
files = (
266+
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
251267
868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
252268
8661386E2CC943DE00F4B78E /* IntegrationTests.swift in Sources */,
253269
);
@@ -405,6 +421,7 @@
405421
MACOSX_DEPLOYMENT_TARGET = 12.0;
406422
MARKETING_VERSION = 1.0;
407423
PRODUCT_BUNDLE_IDENTIFIER = com.google.firebase.VertexAITestApp;
424+
PRODUCT_MODULE_NAME = VertexAITestApp;
408425
PRODUCT_NAME = "$(TARGET_NAME)";
409426
SDKROOT = auto;
410427
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
@@ -441,6 +458,7 @@
441458
MACOSX_DEPLOYMENT_TARGET = 12.0;
442459
MARKETING_VERSION = 1.0;
443460
PRODUCT_BUNDLE_IDENTIFIER = com.google.firebase.VertexAITestApp;
461+
PRODUCT_MODULE_NAME = VertexAITestApp;
444462
PRODUCT_NAME = "$(TARGET_NAME)";
445463
SDKROOT = auto;
446464
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";

0 commit comments

Comments
 (0)