Skip to content

Commit f33b4cb

Browse files
andrewheardandrewheard
authored
[Firebase AI] Add App Check Limited-Use mode integration tests (#15230)
1 parent 4ca80d5 commit f33b4cb

File tree

2 files changed

+73
-28
lines changed

2 files changed

+73
-28
lines changed

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

Lines changed: 33 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -50,13 +50,18 @@ struct GenerateContentIntegrationTests {
5050
@Test(arguments: [
5151
(InstanceConfig.vertexAI_v1beta, ModelNames.gemini2FlashLite),
5252
(InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2FlashLite),
53-
(InstanceConfig.vertexAI_v1beta_staging, ModelNames.gemini2FlashLite),
53+
(InstanceConfig.vertexAI_v1beta_global_appCheckLimitedUse, ModelNames.gemini2FlashLite),
5454
(InstanceConfig.googleAI_v1beta, ModelNames.gemini2FlashLite),
55+
(InstanceConfig.googleAI_v1beta_appCheckLimitedUse, ModelNames.gemini2FlashLite),
5556
(InstanceConfig.googleAI_v1beta, ModelNames.gemma3_4B),
56-
(InstanceConfig.googleAI_v1beta_staging, ModelNames.gemini2FlashLite),
57-
(InstanceConfig.googleAI_v1beta_staging, ModelNames.gemma3_4B),
58-
(InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2FlashLite),
59-
(InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemma3_4B),
57+
(InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemini2FlashLite),
58+
(InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemma3_4B),
59+
// Note: The following configs are commented out for easy one-off manual testing.
60+
// (InstanceConfig.googleAI_v1beta_staging, ModelNames.gemini2FlashLite),
61+
// (InstanceConfig.googleAI_v1beta_staging, ModelNames.gemma3_4B),
62+
// (InstanceConfig.vertexAI_v1beta_staging, ModelNames.gemini2FlashLite),
63+
// (InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2FlashLite),
64+
// (InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemma3_4B),
6065
])
6166
func generateContent(_ config: InstanceConfig, modelName: String) async throws {
6267
let model = FirebaseAI.componentInstance(config).generativeModel(
@@ -165,15 +170,16 @@ struct GenerateContentIntegrationTests {
165170
(.googleAI_v1beta_freeTier, ModelNames.gemini2_5_Flash, ThinkingConfig(
166171
thinkingBudget: 24576, includeThoughts: true
167172
)),
168-
(.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
169-
thinkingBudget: 0
170-
)),
171-
(.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
172-
thinkingBudget: 24576
173-
)),
174-
(.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
175-
thinkingBudget: 24576, includeThoughts: true
176-
)),
173+
// Note: The following configs are commented out for easy one-off manual testing.
174+
// (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
175+
// thinkingBudget: 0
176+
// )),
177+
// (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
178+
// thinkingBudget: 24576
179+
// )),
180+
// (.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2_5_Flash, ThinkingConfig(
181+
// thinkingBudget: 24576, includeThoughts: true
182+
// )),
177183
] as [(InstanceConfig, String, ThinkingConfig)]
178184
)
179185
func generateContentThinking(_ config: InstanceConfig, modelName: String,
@@ -319,8 +325,9 @@ struct GenerateContentIntegrationTests {
319325
InstanceConfig.vertexAI_v1beta,
320326
InstanceConfig.vertexAI_v1beta_global,
321327
InstanceConfig.googleAI_v1beta,
322-
InstanceConfig.googleAI_v1beta_staging,
323-
InstanceConfig.googleAI_v1beta_freeTier_bypassProxy,
328+
// Note: The following configs are commented out for easy one-off manual testing.
329+
// InstanceConfig.googleAI_v1beta_staging,
330+
// InstanceConfig.googleAI_v1beta_freeTier_bypassProxy,
324331
])
325332
func generateImage(_ config: InstanceConfig) async throws {
326333
let generationConfig = GenerationConfig(
@@ -416,13 +423,18 @@ struct GenerateContentIntegrationTests {
416423
@Test(arguments: [
417424
(InstanceConfig.vertexAI_v1beta, ModelNames.gemini2FlashLite),
418425
(InstanceConfig.vertexAI_v1beta_global, ModelNames.gemini2FlashLite),
419-
(InstanceConfig.vertexAI_v1beta_staging, ModelNames.gemini2FlashLite),
426+
(InstanceConfig.vertexAI_v1beta_global_appCheckLimitedUse, ModelNames.gemini2FlashLite),
420427
(InstanceConfig.googleAI_v1beta, ModelNames.gemini2FlashLite),
428+
(InstanceConfig.googleAI_v1beta_appCheckLimitedUse, ModelNames.gemini2FlashLite),
421429
(InstanceConfig.googleAI_v1beta, ModelNames.gemma3_4B),
422-
(InstanceConfig.googleAI_v1beta_staging, ModelNames.gemini2FlashLite),
423-
(InstanceConfig.googleAI_v1beta_staging, ModelNames.gemma3_4B),
424-
(InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2FlashLite),
425-
(InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemma3_4B),
430+
(InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemini2FlashLite),
431+
(InstanceConfig.googleAI_v1beta_freeTier, ModelNames.gemma3_4B),
432+
// Note: The following configs are commented out for easy one-off manual testing.
433+
// (InstanceConfig.vertexAI_v1beta_staging, ModelNames.gemini2FlashLite),
434+
// (InstanceConfig.googleAI_v1beta_staging, ModelNames.gemini2FlashLite),
435+
// (InstanceConfig.googleAI_v1beta_staging, ModelNames.gemma3_4B),
436+
// (InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemini2FlashLite),
437+
// (InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, ModelNames.gemma3_4B),
426438
])
427439
func generateContentStream(_ config: InstanceConfig, modelName: String) async throws {
428440
let expectedResponse = [

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

Lines changed: 40 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -27,12 +27,20 @@ struct InstanceConfig: Equatable, Encodable {
2727
location: "global",
2828
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
2929
)
30+
static let vertexAI_v1beta_global_appCheckLimitedUse = InstanceConfig(
31+
location: "global",
32+
useLimitedUseAppCheckTokens: true,
33+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
34+
)
3035
static let vertexAI_v1beta_staging = InstanceConfig(
3136
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyStaging), version: .v1beta)
3237
)
3338
static let googleAI_v1beta = InstanceConfig(
3439
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
3540
)
41+
static let googleAI_v1beta_appCheckLimitedUse = InstanceConfig(
42+
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
43+
)
3644
static let googleAI_v1beta_staging = InstanceConfig(
3745
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyStaging), version: .v1beta)
3846
)
@@ -48,33 +56,54 @@ struct InstanceConfig: Equatable, Encodable {
4856
static let allConfigs = [
4957
vertexAI_v1beta,
5058
vertexAI_v1beta_global,
51-
vertexAI_v1beta_staging,
59+
vertexAI_v1beta_global_appCheckLimitedUse,
5260
googleAI_v1beta,
53-
googleAI_v1beta_staging,
54-
googleAI_v1beta_freeTier_bypassProxy,
61+
googleAI_v1beta_appCheckLimitedUse,
62+
googleAI_v1beta_freeTier,
63+
// Note: The following configs are commented out for easy one-off manual testing.
64+
// vertexAI_v1beta_staging,
65+
// googleAI_v1beta_staging,
66+
// googleAI_v1beta_freeTier_bypassProxy,
5567
]
5668

5769
static let vertexAI_v1beta_appCheckNotConfigured = InstanceConfig(
5870
appName: FirebaseAppNames.appCheckNotConfigured,
5971
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
6072
)
73+
static let vertexAI_v1beta_appCheckNotConfigured_limitedUseTokens = InstanceConfig(
74+
appName: FirebaseAppNames.appCheckNotConfigured,
75+
useLimitedUseAppCheckTokens: true,
76+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
77+
)
6178
static let googleAI_v1beta_appCheckNotConfigured = InstanceConfig(
6279
appName: FirebaseAppNames.appCheckNotConfigured,
6380
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
6481
)
82+
static let googleAI_v1beta_appCheckNotConfigured_limitedUseTokens = InstanceConfig(
83+
appName: FirebaseAppNames.appCheckNotConfigured,
84+
useLimitedUseAppCheckTokens: true,
85+
apiConfig: APIConfig(service: .googleAI(endpoint: .firebaseProxyProd), version: .v1beta)
86+
)
6587

6688
static let appCheckNotConfiguredConfigs = [
6789
vertexAI_v1beta_appCheckNotConfigured,
90+
vertexAI_v1beta_appCheckNotConfigured_limitedUseTokens,
6891
googleAI_v1beta_appCheckNotConfigured,
92+
googleAI_v1beta_appCheckNotConfigured_limitedUseTokens,
6993
]
7094

7195
let appName: String?
7296
let location: String?
97+
let useLimitedUseAppCheckTokens: Bool
7398
let apiConfig: APIConfig
7499

75-
init(appName: String? = nil, location: String? = nil, apiConfig: APIConfig) {
100+
init(appName: String? = nil,
101+
location: String? = nil,
102+
useLimitedUseAppCheckTokens: Bool = false,
103+
apiConfig: APIConfig) {
76104
self.appName = appName
77105
self.location = location
106+
self.useLimitedUseAppCheckTokens = useLimitedUseAppCheckTokens
78107
self.apiConfig = apiConfig
79108
}
80109

@@ -108,8 +137,12 @@ extension InstanceConfig: CustomTestStringConvertible {
108137
" - Bypass Proxy"
109138
}
110139
let locationSuffix = location.map { " - \($0)" } ?? ""
140+
let appCheckLimitedUseDesignator = useLimitedUseAppCheckTokens ? " - FAC Limited-Use" : ""
111141

112-
return "\(serviceName) (\(versionName))\(freeTierDesignator)\(endpointSuffix)\(locationSuffix)"
142+
return """
143+
\(serviceName) (\(versionName))\(freeTierDesignator)\(endpointSuffix)\(locationSuffix)\
144+
\(appCheckLimitedUseDesignator)
145+
"""
113146
}
114147
}
115148

@@ -122,7 +155,7 @@ extension FirebaseAI {
122155
app: instanceConfig.app,
123156
location: location,
124157
apiConfig: instanceConfig.apiConfig,
125-
useLimitedUseAppCheckTokens: false
158+
useLimitedUseAppCheckTokens: instanceConfig.useLimitedUseAppCheckTokens
126159
)
127160
case .googleAI:
128161
assert(
@@ -133,7 +166,7 @@ extension FirebaseAI {
133166
app: instanceConfig.app,
134167
location: nil,
135168
apiConfig: instanceConfig.apiConfig,
136-
useLimitedUseAppCheckTokens: false
169+
useLimitedUseAppCheckTokens: instanceConfig.useLimitedUseAppCheckTokens
137170
)
138171
}
139172
}

0 commit comments

Comments
 (0)