Skip to content

Commit

Permalink
[Vertex AI] Add integration tests for App Check failure (#14008)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard authored Nov 1, 2024
1 parent 33a44a9 commit 5366717
Show file tree
Hide file tree
Showing 5 changed files with 133 additions and 5 deletions.
2 changes: 1 addition & 1 deletion FirebaseVertexAI/Tests/TestApp/Sources/TestApp.swift
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,7 @@ import SwiftUI
@main
struct TestApp: App {
init() {
AppCheck.setAppCheckProviderFactory(AppCheckDebugProviderFactory())
AppCheck.setAppCheckProviderFactory(TestAppCheckProviderFactory())
FirebaseApp.configure()
}

Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
// Copyright 2024 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 FirebaseAppCheck
import FirebaseCore
import Foundation

/// An `AppCheckProviderFactory` for the Test App.
///
/// Defaults to the `AppCheckDebugProvider` unless the `FirebaseApp` `name` contains
/// ``notConfiguredName``, in which case App Check is not configured; this facilitates integration
/// testing of App Check failure cases.
public class TestAppCheckProviderFactory: NSObject, AppCheckProviderFactory {
/// The name, or a substring of the name, of Firebase apps where App Check is not configured.
public static let notConfiguredName = "app-check-not-configured"

/// Returns the `AppCheckDebugProvider` unless `app.name` contains ``notConfiguredName``.
public func createProvider(with app: FirebaseApp) -> (any AppCheckProvider)? {
if app.name.contains(TestAppCheckProviderFactory.notConfiguredName) {
return nil
}

return AppCheckDebugProvider(app: app)
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -16,9 +16,9 @@ import FirebaseAuth
import FirebaseCore
import FirebaseStorage
import FirebaseVertexAI
import VertexAITestApp
import XCTest

@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
final class IntegrationTests: XCTestCase {
// Set temperature, topP and topK to lowest allowed values to make responses more deterministic.
let generationConfig = GenerationConfig(
Expand Down Expand Up @@ -75,6 +75,21 @@ final class IntegrationTests: XCTestCase {
XCTAssertEqual(text, "Mountain View")
}

func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
addTeardownBlock { await app.delete() }
let vertex = VertexAI.vertexAI(app: app)
let model = vertex.generativeModel(modelName: "gemini-1.5-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 {
Expand Down Expand Up @@ -205,6 +220,21 @@ final class IntegrationTests: XCTestCase {
XCTAssertEqual(response.totalTokens, 34)
XCTAssertEqual(response.totalBillableCharacters, 59)
}

func testCountTokens_appCheckNotConfigured_shouldFail() async throws {
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
addTeardownBlock { await app.delete() }
let vertex = VertexAI.vertexAI(app: app)
let model = vertex.generativeModel(modelName: "gemini-1.5-flash")
let prompt = "Why is the sky blue?"

do {
_ = try await model.countTokens(prompt)
XCTFail("Expected a Firebase App Check error; none thrown.")
} catch {
XCTAssertTrue(String(describing: error).contains("Firebase App Check token is invalid"))
}
}
}

extension StorageReference {
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
// Copyright 2024 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 another `FirebaseApp` with the specified `name` and the same `FirebaseOptions`.
func namedCopy(name: String) throws -> FirebaseApp {
FirebaseApp.configure(name: name, options: options)
guard let app = FirebaseApp.app(name: name) else {
throw AppNotFound(name: name)
}
return app
}

/// Configures an app with the specified `name` and the same `FirebaseOptions` as the default app.
static func defaultNamedCopy(name: String) throws -> FirebaseApp {
guard FirebaseApp.isDefaultAppConfigured(), let defaultApp = FirebaseApp.app() else {
throw DefaultAppNotConfigured()
}
return try defaultApp.namedCopy(name: name)
}

struct AppNotFound: Error {
let name: String
}

struct DefaultAppNotConfigured: Error {}
}
Original file line number Diff line number Diff line change
Expand Up @@ -18,6 +18,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 */; };
/* End PBXBuildFile section */

/* Begin PBXContainerItemProxy section */
Expand All @@ -41,6 +43,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>"; };
/* End PBXFileReference section */

/* Begin PBXFrameworksBuildPhase section */
Expand Down Expand Up @@ -107,6 +111,7 @@
isa = PBXGroup;
children = (
8661385B2CC943DD00F4B78E /* TestApp.swift */,
8698D7472CD4332B00ABA833 /* TestAppCheckProviderFactory.swift */,
8661385D2CC943DD00F4B78E /* ContentView.swift */,
);
path = Sources;
Expand All @@ -125,10 +130,19 @@
isa = PBXGroup;
children = (
868A7C572CCC27AF00E449DD /* Integration */,
8698D7442CD3CEF700ABA833 /* Utilities */,
);
path = Tests;
sourceTree = "<group>";
};
8698D7442CD3CEF700ABA833 /* Utilities */ = {
isa = PBXGroup;
children = (
8698D7452CD3CF2F00ABA833 /* FirebaseAppTestUtils.swift */,
);
path = Utilities;
sourceTree = "<group>";
};
/* End PBXGroup section */

/* Begin PBXNativeTarget section */
Expand Down Expand Up @@ -181,7 +195,7 @@
attributes = {
BuildIndependentTargetsInParallel = 1;
LastSwiftUpdateCheck = 1520;
LastUpgradeCheck = 1520;
LastUpgradeCheck = 1600;
TargetAttributes = {
866138572CC943DD00F4B78E = {
CreatedOnToolsVersion = 15.2;
Expand Down Expand Up @@ -241,13 +255,15 @@
files = (
8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */,
8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */,
8698D7482CD4332B00ABA833 /* TestAppCheckProviderFactory.swift in Sources */,
);
runOnlyForDeploymentPostprocessing = 0;
};
866138652CC943DE00F4B78E /* Sources */ = {
isa = PBXSourcesBuildPhase;
buildActionMask = 2147483647;
files = (
8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
8661386E2CC943DE00F4B78E /* IntegrationTests.swift in Sources */,
);
Expand Down Expand Up @@ -298,6 +314,7 @@
CLANG_WARN_UNREACHABLE_CODE = YES;
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
COPY_PHASE_STRIP = NO;
DEAD_CODE_STRIPPING = YES;
DEBUG_INFORMATION_FORMAT = dwarf;
ENABLE_STRICT_OBJC_MSGSEND = YES;
ENABLE_TESTABILITY = YES;
Expand Down Expand Up @@ -359,6 +376,7 @@
CLANG_WARN_UNREACHABLE_CODE = YES;
CLANG_WARN__DUPLICATE_METHOD_MATCH = YES;
COPY_PHASE_STRIP = NO;
DEAD_CODE_STRIPPING = YES;
DEBUG_INFORMATION_FORMAT = "dwarf-with-dsym";
ENABLE_NS_ASSERTIONS = NO;
ENABLE_STRICT_OBJC_MSGSEND = YES;
Expand Down Expand Up @@ -386,6 +404,7 @@
CODE_SIGN_ENTITLEMENTS = Resources/TestApp.entitlements;
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEAD_CODE_STRIPPING = YES;
DEVELOPMENT_ASSET_PATHS = "\"Resources/Preview Content\"";
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
Expand All @@ -405,6 +424,7 @@
MACOSX_DEPLOYMENT_TARGET = 12.0;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = com.google.firebase.VertexAITestApp;
PRODUCT_MODULE_NAME = VertexAITestApp;
PRODUCT_NAME = "$(TARGET_NAME)";
SDKROOT = auto;
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
Expand All @@ -422,6 +442,7 @@
CODE_SIGN_ENTITLEMENTS = Resources/TestApp.entitlements;
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEAD_CODE_STRIPPING = YES;
DEVELOPMENT_ASSET_PATHS = "\"Resources/Preview Content\"";
ENABLE_PREVIEWS = YES;
GENERATE_INFOPLIST_FILE = YES;
Expand All @@ -441,6 +462,7 @@
MACOSX_DEPLOYMENT_TARGET = 12.0;
MARKETING_VERSION = 1.0;
PRODUCT_BUNDLE_IDENTIFIER = com.google.firebase.VertexAITestApp;
PRODUCT_MODULE_NAME = VertexAITestApp;
PRODUCT_NAME = "$(TARGET_NAME)";
SDKROOT = auto;
SUPPORTED_PLATFORMS = "iphoneos iphonesimulator macosx";
Expand All @@ -453,10 +475,10 @@
866138812CC943DE00F4B78E /* Debug */ = {
isa = XCBuildConfiguration;
buildSettings = {
ALWAYS_EMBED_SWIFT_STANDARD_LIBRARIES = YES;
BUNDLE_LOADER = "$(TEST_HOST)";
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEAD_CODE_STRIPPING = YES;
GENERATE_INFOPLIST_FILE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 15.0;
MACOSX_DEPLOYMENT_TARGET = 12.0;
Expand All @@ -475,10 +497,10 @@
866138822CC943DE00F4B78E /* Release */ = {
isa = XCBuildConfiguration;
buildSettings = {
ALWAYS_EMBED_SWIFT_STANDARD_LIBRARIES = YES;
BUNDLE_LOADER = "$(TEST_HOST)";
CODE_SIGN_STYLE = Automatic;
CURRENT_PROJECT_VERSION = 1;
DEAD_CODE_STRIPPING = YES;
GENERATE_INFOPLIST_FILE = YES;
IPHONEOS_DEPLOYMENT_TARGET = 15.0;
MACOSX_DEPLOYMENT_TARGET = 12.0;
Expand Down

0 comments on commit 5366717

Please sign in to comment.