Skip to content

Commit c4a4550

Browse files
andrewheardandrewheard
authored
[Vertex AI] Add integration tests for staging endpoint (#14637)
1 parent 5c507dd commit c4a4550

File tree

11 files changed

+66
-21
lines changed

11 files changed

+66
-21
lines changed

‎FirebaseVertexAI/Sources/Types/Internal/APIConfig.swift

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -45,7 +45,7 @@ extension APIConfig {
4545
/// See the [Cloud
4646
/// docs](https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/inference) for
4747
/// more details.
48-
case vertexAI
48+
case vertexAI(endpoint: Endpoint)
4949

5050
/// The Gemini Developer API provided by Google AI.
5151
///
@@ -57,8 +57,8 @@ extension APIConfig {
5757
/// This must correspond with the API set in `service`.
5858
var endpoint: Endpoint {
5959
switch self {
60-
case .vertexAI:
61-
return .firebaseVertexAIProd
60+
case let .vertexAI(endpoint: endpoint):
61+
return endpoint
6262
case let .developer(endpoint: endpoint):
6363
return endpoint
6464
}

‎FirebaseVertexAI/Sources/VertexAI.swift

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,9 @@ public class VertexAI {
3838
public static func vertexAI(app: FirebaseApp? = nil,
3939
location: String = "us-central1") -> VertexAI {
4040
let vertexInstance = vertexAI(app: app, location: location, apiConfig: defaultVertexAIAPIConfig)
41-
assert(vertexInstance.apiConfig.service == .vertexAI)
41+
// Verify that the `VertexAI` instance is always configured with the production endpoint since
42+
// this is the public API surface for creating an instance.
43+
assert(vertexInstance.apiConfig.service == .vertexAI(endpoint: .firebaseVertexAIProd))
4244
assert(vertexInstance.apiConfig.service.endpoint == .firebaseVertexAIProd)
4345
assert(vertexInstance.apiConfig.version == .v1beta)
4446

@@ -155,7 +157,10 @@ public class VertexAI {
155157

156158
let location: String?
157159

158-
static let defaultVertexAIAPIConfig = APIConfig(service: .vertexAI, version: .v1beta)
160+
static let defaultVertexAIAPIConfig = APIConfig(
161+
service: .vertexAI(endpoint: .firebaseVertexAIProd),
162+
version: .v1beta
163+
)
159164

160165
static func vertexAI(app: FirebaseApp?, location: String?, apiConfig: APIConfig) -> VertexAI {
161166
guard let app = app ?? FirebaseApp.app() else {

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,7 +71,9 @@ struct CountTokensIntegrationTests {
7171

7272
@Test(arguments: [
7373
InstanceConfig.vertexV1,
74+
InstanceConfig.vertexV1Staging,
7475
InstanceConfig.vertexV1Beta,
76+
InstanceConfig.vertexV1BetaStaging,
7577
/* System instructions are not supported on the v1 Developer API. */
7678
InstanceConfig.developerV1Beta,
7779
])

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

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -78,7 +78,9 @@ struct GenerateContentIntegrationTests {
7878
"Generate an enum and provide a system instruction",
7979
arguments: [
8080
InstanceConfig.vertexV1,
81+
InstanceConfig.vertexV1Staging,
8182
InstanceConfig.vertexV1Beta,
83+
InstanceConfig.vertexV1BetaStaging,
8284
/* System instructions are not supported on the v1 Developer API. */
8385
InstanceConfig.developerV1Beta,
8486
]

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

Lines changed: 20 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,17 @@ import VertexAITestApp
2020
@testable import class FirebaseVertexAI.VertexAI
2121

2222
struct InstanceConfig {
23-
static let vertexV1 = InstanceConfig(apiConfig: APIConfig(service: .vertexAI, version: .v1))
23+
static let vertexV1 = InstanceConfig(
24+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1)
25+
)
26+
static let vertexV1Staging = InstanceConfig(
27+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIStaging), version: .v1)
28+
)
2429
static let vertexV1Beta = InstanceConfig(
25-
apiConfig: APIConfig(service: .vertexAI, version: .v1beta)
30+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1beta)
31+
)
32+
static let vertexV1BetaStaging = InstanceConfig(
33+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIStaging), version: .v1beta)
2634
)
2735
static let developerV1 = InstanceConfig(
2836
appName: FirebaseAppNames.spark,
@@ -32,15 +40,22 @@ struct InstanceConfig {
3240
appName: FirebaseAppNames.spark,
3341
apiConfig: APIConfig(service: .developer(endpoint: .generativeLanguage), version: .v1beta)
3442
)
35-
static let allConfigs = [vertexV1, vertexV1Beta, developerV1, developerV1Beta]
43+
static let allConfigs = [
44+
vertexV1,
45+
vertexV1Staging,
46+
vertexV1Beta,
47+
vertexV1BetaStaging,
48+
developerV1,
49+
developerV1Beta,
50+
]
3651

3752
static let vertexV1AppCheckNotConfigured = InstanceConfig(
3853
appName: FirebaseAppNames.appCheckNotConfigured,
39-
apiConfig: APIConfig(service: .vertexAI, version: .v1)
54+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1)
4055
)
4156
static let vertexV1BetaAppCheckNotConfigured = InstanceConfig(
4257
appName: FirebaseAppNames.appCheckNotConfigured,
43-
apiConfig: APIConfig(service: .vertexAI, version: .v1beta)
58+
apiConfig: APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1beta)
4459
)
4560

4661
let appName: String?

‎FirebaseVertexAI/Tests/Unit/ChatTests.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,7 @@ final class ChatTests: XCTestCase {
6666
firebaseAppID: "My app ID",
6767
firebaseApp: app
6868
),
69-
apiConfig: APIConfig(service: .vertexAI, version: .v1beta),
69+
apiConfig: VertexAI.defaultVertexAIAPIConfig,
7070
tools: nil,
7171
requestOptions: RequestOptions(),
7272
urlSession: urlSession

‎FirebaseVertexAI/Tests/Unit/GenerativeModelTests.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -58,7 +58,7 @@ final class GenerativeModelTests: XCTestCase {
5858
].sorted()
5959
let testModelResourceName =
6060
"projects/test-project-id/locations/test-location/publishers/google/models/test-model"
61-
let apiConfig = APIConfig(service: .vertexAI, version: .v1beta)
61+
let apiConfig = VertexAI.defaultVertexAIAPIConfig
6262

6363
let vertexSubdirectory = "vertexai"
6464

‎FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenGenerationRequestTests.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,7 @@ final class ImagenGenerationRequestTests: XCTestCase {
3636
addWatermark: nil,
3737
includeResponsibleAIFilterReason: includeResponsibleAIFilterReason
3838
)
39-
let apiConfig = APIConfig(service: .vertexAI, version: .v1beta)
39+
let apiConfig = VertexAI.defaultVertexAIAPIConfig
4040

4141
let instance = ImageGenerationInstance(prompt: "test-prompt")
4242

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

Lines changed: 23 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,19 +19,40 @@ import XCTest
1919
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
2020
final class APIConfigTests: XCTestCase {
2121
func testInitialize_vertexAI_prod_v1() {
22-
let apiConfig = APIConfig(service: .vertexAI, version: .v1)
22+
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1)
2323

2424
XCTAssertEqual(apiConfig.service.endpoint.rawValue, "https://firebasevertexai.googleapis.com")
2525
XCTAssertEqual(apiConfig.version.rawValue, "v1")
2626
}
2727

2828
func testInitialize_vertexAI_prod_v1beta() {
29-
let apiConfig = APIConfig(service: .vertexAI, version: .v1beta)
29+
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIProd), version: .v1beta)
3030

3131
XCTAssertEqual(apiConfig.service.endpoint.rawValue, "https://firebasevertexai.googleapis.com")
3232
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
3333
}
3434

35+
func testInitialize_vertexAI_staging_v1() {
36+
let apiConfig = APIConfig(service: .vertexAI(endpoint: .firebaseVertexAIStaging), version: .v1)
37+
38+
XCTAssertEqual(
39+
apiConfig.service.endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com"
40+
)
41+
XCTAssertEqual(apiConfig.version.rawValue, "v1")
42+
}
43+
44+
func testInitialize_vertexAI_staging_v1beta() {
45+
let apiConfig = APIConfig(
46+
service: .vertexAI(endpoint: .firebaseVertexAIStaging),
47+
version: .v1beta
48+
)
49+
50+
XCTAssertEqual(
51+
apiConfig.service.endpoint.rawValue, "https://staging-firebasevertexai.sandbox.googleapis.com"
52+
)
53+
XCTAssertEqual(apiConfig.version.rawValue, "v1beta")
54+
}
55+
3556
func testInitialize_developer_staging_v1beta() {
3657
let apiConfig = APIConfig(
3758
service: .developer(endpoint: .firebaseVertexAIStaging), version: .v1beta

‎FirebaseVertexAI/Tests/Unit/Types/Internal/Requests/CountTokensRequestTests.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ final class CountTokensRequestTests: XCTestCase {
2323

2424
let modelResourceName = "models/test-model-name"
2525
let textPart = TextPart("test-prompt")
26-
let vertexAPIConfig = APIConfig(service: .vertexAI, version: .v1beta)
26+
let vertexAPIConfig = VertexAI.defaultVertexAIAPIConfig
2727
let developerAPIConfig = APIConfig(
2828
service: .developer(endpoint: .firebaseVertexAIProd),
2929
version: .v1beta

0 commit comments

Comments
 (0)