Skip to content

Commit e04e66a

Browse files
andrewheardandrewheard
committed
Fix unit tests
1 parent 50cee21 commit e04e66a

File tree

3 files changed

+104
-51
lines changed

3 files changed

+104
-51
lines changed

‎FirebaseAI/Tests/Unit/Types/BackendTests.swift‎

Lines changed: 3 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,27 +19,25 @@ import XCTest
1919
final class BackendTests: XCTestCase {
2020
func testVertexAI_defaultLocation() {
2121
let expectedAPIConfig = APIConfig(
22-
service: .vertexAI(endpoint: .firebaseProxyProd),
22+
service: .vertexAI(endpoint: .firebaseProxyProd, location: "us-central1"),
2323
version: .v1beta
2424
)
2525

2626
let backend = Backend.vertexAI()
2727

2828
XCTAssertEqual(backend.apiConfig, expectedAPIConfig)
29-
XCTAssertEqual(backend.location, "us-central1")
3029
}
3130

3231
func testVertexAI_customLocation() {
32+
let customLocation = "europe-west1"
3333
let expectedAPIConfig = APIConfig(
34-
service: .vertexAI(endpoint: .firebaseProxyProd),
34+
service: .vertexAI(endpoint: .firebaseProxyProd, location: customLocation),
3535
version: .v1beta
3636
)
37-
let customLocation = "europe-west1"
3837

3938
let backend = Backend.vertexAI(location: customLocation)
4039

4140
XCTAssertEqual(backend.apiConfig, expectedAPIConfig)
42-
XCTAssertEqual(backend.location, customLocation)
4341
}
4442

4543
func testGoogleAI() {
@@ -51,6 +49,5 @@ final class BackendTests: XCTestCase {
5149
let backend = Backend.googleAI()
5250

5351
XCTAssertEqual(backend.apiConfig, expectedAPIConfig)
54-
XCTAssertNil(backend.location)
5552
}
5653
}

‎FirebaseAI/Tests/Unit/Types/Internal/APIConfigTests.swift‎

Lines changed: 56 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -18,38 +18,70 @@ import XCTest
1818

1919
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2020
final class APIConfigTests: XCTestCase {
21+
let defaultLocation = "us-central1"
22+
let globalLocation = "global"
23+
2124
func testInitialize_vertexAI_prod_v1() {
22-
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1)
25+
let apiConfig = APIConfig(
26+
service: .vertexAI(endpoint: .firebaseProxyProd, location: defaultLocation),
27+
version: .v1
28+
)
2329

24-
XCTAssertEqual(apiConfig.service.endpoint.rawValue, "https://firebasevertexai.googleapis.com")
30+
switch apiConfig.service {
31+
case let .vertexAI(endpoint: endpoint, location: location):
32+
XCTAssertEqual(endpoint.rawValue, "https://firebasevertexai.googleapis.com")
33+
XCTAssertEqual(location, defaultLocation)
34+
case .googleAI:
35+
XCTFail("Expected .vertexAI, got .googleAI")
36+
}
2537
XCTAssertEqual(apiConfig.version.rawValue, "v1")
2638
}
2739

2840
func testInitialize_vertexAI_prod_v1beta() {
29-
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta)
41+
let apiConfig = APIConfig(
42+
service: .vertexAI(endpoint: .firebaseProxyProd, location: defaultLocation),
43+
version: .v1beta
44+
)
3045

31-
XCTAssertEqual(apiConfig.service.endpoint.rawValue, "https://firebasevertexai.googleapis.com")
46+
switch apiConfig.service {
47+
case let .vertexAI(endpoint: endpoint, location: location):
48+
XCTAssertEqual(endpoint.rawValue, "https://firebasevertexai.googleapis.com")
49+
XCTAssertEqual(location, defaultLocation)
50+
case .googleAI:
51+
XCTFail("Expected .vertexAI, got .googleAI")
52+
}
3253
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
3354
}
3455

3556
func testInitialize_vertexAI_staging_v1() {
36-
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseProxyStaging), version: .v1)
37-
38-
XCTAssertEqual(
39-
apiConfig.service.endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com"
57+
let apiConfig = APIConfig(
58+
service: .vertexAI(endpoint: .firebaseProxyStaging, location: defaultLocation),
59+
version: .v1
4060
)
61+
62+
switch apiConfig.service {
63+
case let .vertexAI(endpoint: endpoint, location: location):
64+
XCTAssertEqual(endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com")
65+
XCTAssertEqual(location, defaultLocation)
66+
case .googleAI:
67+
XCTFail("Expected .vertexAI, got .googleAI")
68+
}
4169
XCTAssertEqual(apiConfig.version.rawValue, "v1")
4270
}
4371

4472
func testInitialize_vertexAI_staging_v1beta() {
4573
let apiConfig = APIConfig(
46-
service: .vertexAI(endpoint: .firebaseProxyStaging),
74+
service: .vertexAI(endpoint: .firebaseProxyStaging, location: defaultLocation),
4775
version: .v1beta
4876
)
4977

50-
XCTAssertEqual(
51-
apiConfig.service.endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com"
52-
)
78+
switch apiConfig.service {
79+
case let .vertexAI(endpoint: endpoint, location: location):
80+
XCTAssertEqual(endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com")
81+
XCTAssertEqual(location, defaultLocation)
82+
case .googleAI:
83+
XCTFail("Expected .vertexAI, got .googleAI")
84+
}
5385
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
5486
}
5587

@@ -58,16 +90,24 @@ final class APIConfigTests: XCTestCase {
5890
service: .googleAI(endpoint: .firebaseProxyStaging), version: .v1beta
5991
)
6092

61-
XCTAssertEqual(
62-
apiConfig.service.endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com"
63-
)
93+
switch apiConfig.service {
94+
case .vertexAI:
95+
XCTFail("Expected .googleAI, got .vertexAI")
96+
case let .googleAI(endpoint: endpoint):
97+
XCTAssertEqual(endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com")
98+
}
6499
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
65100
}
66101

67102
func testInitialize_developer_generativeLanguage_v1beta() {
68103
let apiConfig = APIConfig(service: .googleAI(endpoint: .googleAIBypassProxy), version: .v1beta)
69104

70-
XCTAssertEqual(apiConfig.service.endpoint.rawValue, "https://generativelanguage.googleapis.com")
105+
switch apiConfig.service {
106+
case .vertexAI:
107+
XCTFail("Expected .googleAI, got .vertexAI")
108+
case let .googleAI(endpoint: endpoint):
109+
XCTAssertEqual(endpoint.rawValue, "https://generativelanguage.googleapis.com")
110+
}
71111
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
72112
}
73113
}

‎FirebaseAI/Tests/Unit/VertexComponentTests.swift‎

Lines changed: 45 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -57,8 +57,9 @@ class VertexComponentTests: XCTestCase {
5757
XCTAssertNotNil(vertex)
5858
XCTAssertEqual(vertex.firebaseInfo.projectID, VertexComponentTests.projectID)
5959
XCTAssertEqual(vertex.firebaseInfo.apiKey, VertexComponentTests.apiKey)
60-
XCTAssertEqual(vertex.location, "us-central1")
61-
XCTAssertEqual(vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd))
60+
XCTAssertEqual(
61+
vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd, location: "us-central1")
62+
)
6263
XCTAssertEqual(vertex.apiConfig.service.endpoint, .firebaseProxyProd)
6364
XCTAssertEqual(vertex.apiConfig.version, .v1beta)
6465
}
@@ -71,8 +72,9 @@ class VertexComponentTests: XCTestCase {
7172
XCTAssertNotNil(vertex)
7273
XCTAssertEqual(vertex.firebaseInfo.projectID, VertexComponentTests.projectID)
7374
XCTAssertEqual(vertex.firebaseInfo.apiKey, VertexComponentTests.apiKey)
74-
XCTAssertEqual(vertex.location, location)
75-
XCTAssertEqual(vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd))
75+
XCTAssertEqual(
76+
vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd, location: location)
77+
)
7678
XCTAssertEqual(vertex.apiConfig.service.endpoint, .firebaseProxyProd)
7779
XCTAssertEqual(vertex.apiConfig.version, .v1beta)
7880
}
@@ -87,8 +89,9 @@ class VertexComponentTests: XCTestCase {
8789
XCTAssertNotNil(vertex)
8890
XCTAssertEqual(vertex.firebaseInfo.projectID, VertexComponentTests.projectID)
8991
XCTAssertEqual(vertex.firebaseInfo.apiKey, VertexComponentTests.apiKey)
90-
XCTAssertEqual(vertex.location, location)
91-
XCTAssertEqual(vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd))
92+
XCTAssertEqual(
93+
vertex.apiConfig.service, .vertexAI(endpoint: .firebaseProxyProd, location: location)
94+
)
9295
XCTAssertEqual(vertex.apiConfig.service.endpoint, .firebaseProxyProd)
9396
XCTAssertEqual(vertex.apiConfig.version, .v1beta)
9497
}
@@ -154,14 +157,17 @@ class VertexComponentTests: XCTestCase {
154157
func testSameAppAndDifferentAPI_newInstanceCreated() throws {
155158
let vertex1 = FirebaseAI.createInstance(
156159
app: VertexComponentTests.app,
157-
location: location,
158-
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1beta),
160+
apiConfig: APIConfig(
161+
service: .vertexAI(endpoint: .firebaseProxyProd, location: location),
162+
version: .v1beta
163+
),
159164
useLimitedUseAppCheckTokens: false
160165
)
161166
let vertex2 = FirebaseAI.createInstance(
162167
app: VertexComponentTests.app,
163-
location: location,
164-
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseProxyProd), version: .v1),
168+
apiConfig: APIConfig(
169+
service: .vertexAI(endpoint: .firebaseProxyProd, location: location), version: .v1
170+
),
165171
useLimitedUseAppCheckTokens: false
166172
)
167173

@@ -182,8 +188,10 @@ class VertexComponentTests: XCTestCase {
182188
weakApp = try XCTUnwrap(app1)
183189
let vertex = FirebaseAI(
184190
app: app1,
185-
location: "transitory location",
186-
apiConfig: FirebaseAI.defaultVertexAIAPIConfig,
191+
apiConfig: APIConfig(
192+
service: .vertexAI(endpoint: .firebaseProxyProd, location: "transitory location"),
193+
version: .v1beta,
194+
),
187195
useLimitedUseAppCheckTokens: false
188196
)
189197
weakVertex = vertex
@@ -195,13 +203,13 @@ class VertexComponentTests: XCTestCase {
195203

196204
func testModelResourceName_vertexAI() throws {
197205
let app = try XCTUnwrap(VertexComponentTests.app)
206+
let location = "test-location"
198207
let vertex = FirebaseAI.firebaseAI(app: app, backend: .vertexAI(location: location))
199208
let model = "test-model-name"
200209
let projectID = vertex.firebaseInfo.projectID
201210

202211
let modelResourceName = vertex.modelResourceName(modelName: model)
203212

204-
let location = try XCTUnwrap(vertex.location)
205213
XCTAssertEqual(
206214
modelResourceName,
207215
"projects/\(projectID)/locations/\(location)/publishers/google/models/\(model)"
@@ -212,10 +220,7 @@ class VertexComponentTests: XCTestCase {
212220
let app = try XCTUnwrap(VertexComponentTests.app)
213221
let apiConfig = APIConfig(service: .googleAI(endpoint: .googleAIBypassProxy), version: .v1beta)
214222
let vertex = FirebaseAI.createInstance(
215-
app: app,
216-
location: nil,
217-
apiConfig: apiConfig,
218-
useLimitedUseAppCheckTokens: false
223+
app: app, apiConfig: apiConfig, useLimitedUseAppCheckTokens: false
219224
)
220225
let model = "test-model-name"
221226

@@ -231,10 +236,7 @@ class VertexComponentTests: XCTestCase {
231236
version: .v1beta
232237
)
233238
let vertex = FirebaseAI.createInstance(
234-
app: app,
235-
location: nil,
236-
apiConfig: apiConfig,
237-
useLimitedUseAppCheckTokens: false
239+
app: app, apiConfig: apiConfig, useLimitedUseAppCheckTokens: false
238240
)
239241
let model = "test-model-name"
240242
let projectID = vertex.firebaseInfo.projectID
@@ -244,33 +246,47 @@ class VertexComponentTests: XCTestCase {
244246
XCTAssertEqual(modelResourceName, "projects/\(projectID)/models/\(model)")
245247
}
246248

247-
func testGenerativeModel_vertexAI() async throws {
249+
func testGenerativeModel_vertexAI_defaultLocation() async throws {
248250
let app = try XCTUnwrap(VertexComponentTests.app)
249-
let vertex = FirebaseAI.firebaseAI(app: app, backend: .vertexAI(location: location))
251+
let vertex = FirebaseAI.firebaseAI(app: app, backend: .vertexAI())
250252
let modelResourceName = vertex.modelResourceName(modelName: modelName)
251253
let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts)
252254

253255
let generativeModel = vertex.generativeModel(
254-
modelName: modelName,
255-
systemInstruction: systemInstruction
256+
modelName: modelName, systemInstruction: systemInstruction
256257
)
257258

258259
XCTAssertEqual(generativeModel.modelResourceName, modelResourceName)
259260
XCTAssertEqual(generativeModel.systemInstruction, expectedSystemInstruction)
260261
XCTAssertEqual(generativeModel.apiConfig, FirebaseAI.defaultVertexAIAPIConfig)
261262
}
262263

264+
func testGenerativeModel_vertexAI_customLocation() async throws {
265+
let app = try XCTUnwrap(VertexComponentTests.app)
266+
let vertex = FirebaseAI.firebaseAI(app: app, backend: .vertexAI(location: location))
267+
let modelResourceName = vertex.modelResourceName(modelName: modelName)
268+
let expectedAPIConfig = APIConfig(
269+
service: .vertexAI(endpoint: .firebaseProxyProd, location: location), version: .v1beta
270+
)
271+
let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts)
272+
273+
let generativeModel = vertex.generativeModel(
274+
modelName: modelName, systemInstruction: systemInstruction
275+
)
276+
277+
XCTAssertEqual(generativeModel.modelResourceName, modelResourceName)
278+
XCTAssertEqual(generativeModel.systemInstruction, expectedSystemInstruction)
279+
XCTAssertEqual(generativeModel.apiConfig, expectedAPIConfig)
280+
}
281+
263282
func testGenerativeModel_developerAPI() async throws {
264283
let app = try XCTUnwrap(VertexComponentTests.app)
265284
let apiConfig = APIConfig(
266285
service: .googleAI(endpoint: .firebaseProxyStaging),
267286
version: .v1beta
268287
)
269288
let vertex = FirebaseAI.createInstance(
270-
app: app,
271-
location: nil,
272-
apiConfig: apiConfig,
273-
useLimitedUseAppCheckTokens: false
289+
app: app, apiConfig: apiConfig, useLimitedUseAppCheckTokens: false
274290
)
275291
let modelResourceName = vertex.modelResourceName(modelName: modelName)
276292
let expectedSystemInstruction = ModelContent(role: nil, parts: systemInstruction.parts)

0 commit comments

Comments
 (0)