Skip to content

Commit c38263e

Browse files
cynthiajoancynthiajoan
authored
Add appId for in request header (#17278)
1 parent 2a0bdc6 commit c38263e

File tree

5 files changed

+230
-5
lines changed

5 files changed

+230
-5
lines changed

‎packages/firebase_vertexai/firebase_vertexai/lib/src/base_model.dart

+4-1
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ abstract class BaseModel {
9494

9595
/// Returns a function that generates Firebase auth tokens.
9696
static FutureOr<Map<String, String>> Function() firebaseTokens(
97-
FirebaseAppCheck? appCheck, FirebaseAuth? auth) {
97+
FirebaseAppCheck? appCheck, FirebaseAuth? auth, FirebaseApp? app) {
9898
return () async {
9999
Map<String, String> headers = {};
100100
// Override the client name in Google AI SDK
@@ -112,6 +112,9 @@ abstract class BaseModel {
112112
headers['Authorization'] = 'Firebase $idToken';
113113
}
114114
}
115+
if (app != null && app.isAutomaticDataCollectionEnabled) {
116+
headers['X-Firebase-AppId'] = app.options.appId;
117+
}
115118
return headers;
116119
};
117120
}

‎packages/firebase_vertexai/firebase_vertexai/lib/src/generative_model.dart

+3-2
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ final class GenerativeModel extends BaseApiClientModel {
5656
client: HttpApiClient(
5757
apiKey: app.options.apiKey,
5858
httpClient: httpClient,
59-
requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
59+
requestHeaders: BaseModel.firebaseTokens(appCheck, auth, app)));
6060

6161
GenerativeModel._constructTestModel({
6262
required String model,
@@ -82,7 +82,8 @@ final class GenerativeModel extends BaseApiClientModel {
8282
client: apiClient ??
8383
HttpApiClient(
8484
apiKey: app.options.apiKey,
85-
requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
85+
requestHeaders:
86+
BaseModel.firebaseTokens(appCheck, auth, app)));
8687

8788
final List<SafetySetting> _safetySettings;
8889
final GenerationConfig? _generationConfig;

‎packages/firebase_vertexai/firebase_vertexai/lib/src/imagen_model.dart

+1-1
Original file line numberDiff line numberDiff line change
@@ -41,7 +41,7 @@ final class ImagenModel extends BaseApiClientModel {
4141
location: location,
4242
client: HttpApiClient(
4343
apiKey: app.options.apiKey,
44-
requestHeaders: BaseModel.firebaseTokens(appCheck, auth)));
44+
requestHeaders: BaseModel.firebaseTokens(appCheck, auth, app)));
4545

4646
final ImagenGenerationConfig? _generationConfig;
4747
final ImagenSafetySettings? _safetySettings;

‎packages/firebase_vertexai/firebase_vertexai/lib/src/live_model.dart

+1-1
Original file line numberDiff line numberDiff line change
@@ -82,7 +82,7 @@ final class LiveGenerativeModel extends BaseModel {
8282
};
8383

8484
final request = jsonEncode(setupJson);
85-
final headers = await BaseModel.firebaseTokens(_appCheck, _auth)();
85+
final headers = await BaseModel.firebaseTokens(_appCheck, _auth, _app)();
8686
var ws = IOWebSocketChannel.connect(Uri.parse(uri), headers: headers);
8787
await ws.ready;
8888

‎packages/firebase_vertexai/firebase_vertexai/test/base_model_test.dart

+221
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,221 @@
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+
//import 'dart:';
15+
import 'package:firebase_app_check/firebase_app_check.dart';
16+
import 'package:firebase_auth/firebase_auth.dart';
17+
import 'package:firebase_core/firebase_core.dart';
18+
import 'package:firebase_vertexai/src/base_model.dart';
19+
import 'package:firebase_vertexai/src/client.dart';
20+
import 'package:flutter_test/flutter_test.dart';
21+
import 'package:mockito/mockito.dart';
22+
23+
// Mock FirebaseApp
24+
// ignore: avoid_implementing_value_types
25+
class MockFirebaseApp extends Mock implements FirebaseApp {
26+
@override
27+
FirebaseOptions get options => MockFirebaseOptions();
28+
29+
@override
30+
bool get isAutomaticDataCollectionEnabled => true;
31+
}
32+
33+
// Mock FirebaseOptions
34+
// ignore: must_be_immutable, avoid_implementing_value_types
35+
class MockFirebaseOptions extends Mock implements FirebaseOptions {
36+
@override
37+
String get projectId => 'test-project';
38+
39+
@override
40+
String get appId => 'test-app-id';
41+
}
42+
43+
// Mock Firebase App Check
44+
class MockFirebaseAppCheck extends Mock implements FirebaseAppCheck {
45+
@override
46+
Future<String?> getToken([bool? forceRefresh = false]) async =>
47+
super.noSuchMethod(Invocation.method(#getToken, [forceRefresh]));
48+
}
49+
50+
// Mock Firebase Auth
51+
class MockFirebaseAuth extends Mock implements FirebaseAuth {
52+
@override
53+
User? get currentUser => super.noSuchMethod(Invocation.getter(#currentUser));
54+
}
55+
56+
// Mock Firebase User
57+
class MockUser extends Mock implements User {
58+
@override
59+
Future<String?> getIdToken([bool? forceRefresh = false]) async =>
60+
super.noSuchMethod(Invocation.method(#getIdToken, [forceRefresh]));
61+
}
62+
63+
class MockApiClient extends Mock implements ApiClient {
64+
@override
65+
Future<Map<String, Object?>> makeRequest(
66+
Uri uri, Map<String, Object?> params) async {
67+
// Simulate a successful API response
68+
return {'mockResponse': 'success'};
69+
}
70+
}
71+
72+
// A concrete subclass of BaseModel for testing purposes
73+
class TestBaseModel extends BaseModel {
74+
TestBaseModel({
75+
required String model,
76+
required String location,
77+
required FirebaseApp app,
78+
}) : super(model: model, location: location, app: app);
79+
}
80+
81+
class TestApiClientModel extends BaseApiClientModel {
82+
TestApiClientModel({
83+
required super.model,
84+
required super.location,
85+
required super.app,
86+
required ApiClient client,
87+
}) : super(client: client);
88+
}
89+
90+
void main() {
91+
group('BaseModel', () {
92+
test('normalizeModelName returns correct prefix and name for model code',
93+
() {
94+
final result = BaseModel.normalizeModelName('models/my-model');
95+
expect(result.prefix, 'models');
96+
expect(result.name, 'my-model');
97+
});
98+
99+
test(
100+
'normalizeModelName returns correct prefix and name for user-friendly name',
101+
() {
102+
final result = BaseModel.normalizeModelName('my-model');
103+
expect(result.prefix, 'models');
104+
expect(result.name, 'my-model');
105+
});
106+
107+
test('taskUri constructs the correct URI for a task', () {
108+
final mockApp = MockFirebaseApp();
109+
final model = TestBaseModel(
110+
model: 'my-model', location: 'us-central1', app: mockApp);
111+
final taskUri = model.taskUri(Task.generateContent);
112+
expect(taskUri.toString(),
113+
'https://firebasevertexai.googleapis.com/v1beta/projects/test-project/locations/us-central1/publishers/google/models/my-model:generateContent');
114+
});
115+
116+
test('taskUri constructs the correct URI for a task with model code', () {
117+
final mockApp = MockFirebaseApp();
118+
final model = TestBaseModel(
119+
model: 'models/my-model', location: 'us-central1', app: mockApp);
120+
final taskUri = model.taskUri(Task.countTokens);
121+
expect(taskUri.toString(),
122+
'https://firebasevertexai.googleapis.com/v1beta/projects/test-project/locations/us-central1/publishers/google/models/my-model:countTokens');
123+
});
124+
125+
test('firebaseTokens returns a function that generates headers', () async {
126+
final tokenFunction = BaseModel.firebaseTokens(null, null, null);
127+
final headers = await tokenFunction();
128+
expect(headers['x-goog-api-client'], contains('gl-dart'));
129+
expect(headers['x-goog-api-client'], contains('fire'));
130+
expect(headers.length, 1);
131+
});
132+
133+
test('firebaseTokens includes App Check token if available', () async {
134+
final mockAppCheck = MockFirebaseAppCheck();
135+
when(mockAppCheck.getToken())
136+
.thenAnswer((_) async => 'test-app-check-token');
137+
final tokenFunction = BaseModel.firebaseTokens(mockAppCheck, null, null);
138+
final headers = await tokenFunction();
139+
expect(headers['X-Firebase-AppCheck'], 'test-app-check-token');
140+
expect(headers['x-goog-api-client'], contains('gl-dart'));
141+
expect(headers['x-goog-api-client'], contains('fire'));
142+
expect(headers.length, 2);
143+
});
144+
145+
test('firebaseTokens includes Auth ID token if available', () async {
146+
final mockAuth = MockFirebaseAuth();
147+
final mockUser = MockUser();
148+
when(mockUser.getIdToken()).thenAnswer((_) async => 'test-id-token');
149+
when(mockAuth.currentUser).thenReturn(mockUser);
150+
final tokenFunction = BaseModel.firebaseTokens(null, mockAuth, null);
151+
final headers = await tokenFunction();
152+
expect(headers['Authorization'], 'Firebase test-id-token');
153+
expect(headers['x-goog-api-client'], contains('gl-dart'));
154+
expect(headers['x-goog-api-client'], contains('fire'));
155+
expect(headers.length, 2);
156+
});
157+
158+
test(
159+
'firebaseTokens includes App ID if automatic data collection is enabled',
160+
() async {
161+
final mockApp = MockFirebaseApp();
162+
163+
final tokenFunction = BaseModel.firebaseTokens(null, null, mockApp);
164+
final headers = await tokenFunction();
165+
expect(headers['X-Firebase-AppId'], 'test-app-id');
166+
expect(headers['x-goog-api-client'], contains('gl-dart'));
167+
expect(headers['x-goog-api-client'], contains('fire'));
168+
expect(headers.length, 2);
169+
});
170+
171+
test('firebaseTokens includes all tokens if available', () async {
172+
final mockAppCheck = MockFirebaseAppCheck();
173+
when(mockAppCheck.getToken())
174+
.thenAnswer((_) async => 'test-app-check-token');
175+
final mockAuth = MockFirebaseAuth();
176+
final mockUser = MockUser();
177+
when(mockUser.getIdToken()).thenAnswer((_) async => 'test-id-token');
178+
when(mockAuth.currentUser).thenReturn(mockUser);
179+
final mockApp = MockFirebaseApp();
180+
181+
final tokenFunction =
182+
BaseModel.firebaseTokens(mockAppCheck, mockAuth, mockApp);
183+
final headers = await tokenFunction();
184+
expect(headers['X-Firebase-AppCheck'], 'test-app-check-token');
185+
expect(headers['Authorization'], 'Firebase test-id-token');
186+
expect(headers['X-Firebase-AppId'], 'test-app-id');
187+
expect(headers['x-goog-api-client'], contains('gl-dart'));
188+
expect(headers['x-goog-api-client'], contains('fire'));
189+
expect(headers.length, 4);
190+
});
191+
});
192+
193+
group('BaseApiClientModel', () {
194+
test('makeRequest returns the parsed response', () async {
195+
final mockApp = MockFirebaseApp();
196+
final mockClient = MockApiClient();
197+
final model = TestApiClientModel(
198+
model: 'test-model',
199+
location: 'us-central1',
200+
app: mockApp,
201+
client: mockClient);
202+
final params = {'input': 'test'};
203+
const task = Task.generateContent;
204+
205+
final response = await model.makeRequest(
206+
task, params, (data) => data['mockResponse']! as String);
207+
expect(response, 'success');
208+
});
209+
210+
test('client getter returns the injected ApiClient', () {
211+
final mockApp = MockFirebaseApp();
212+
final mockClient = MockApiClient();
213+
final model = TestApiClientModel(
214+
model: 'test-model',
215+
location: 'us-central1',
216+
app: mockApp,
217+
client: mockClient);
218+
expect(model.client, mockClient);
219+
});
220+
});
221+
}

0 commit comments

Comments
 (0)