Skip to content

Commit a89956f

Browse files
committed
VertexAI Component tests; fix instance generation for different regions (#12631)
1 parent 4e3a4d1 commit a89956f

File tree

4 files changed

+132
-8
lines changed

4 files changed

+132
-8
lines changed

FirebaseVertexAI/Sources/VertexAI.swift

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -94,7 +94,7 @@ public class VertexAI: NSObject {
9494

9595
private let appCheck: AppCheckInterop?
9696

97-
private let region: String
97+
let region: String
9898

9999
init(app: FirebaseApp, region: String) {
100100
self.app = app

FirebaseVertexAI/Sources/VertexAIComponent.swift

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,7 @@ import Foundation
2222
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
2323
@objc(FIRVertexAIProvider)
2424
protocol VertexAIProvider {
25-
@objc func vertexAI(_ location: String) -> VertexAI
25+
@objc func vertexAI(_ region: String) -> VertexAI
2626
}
2727

2828
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
@@ -32,7 +32,7 @@ class VertexAIComponent: NSObject, Library, VertexAIProvider {
3232

3333
/// The app associated with all `VertexAI` instances in this container.
3434
/// This is `unowned` instead of `weak` so it can be used without unwrapping in `vertexAI(...)`
35-
private unowned let app: FirebaseApp
35+
unowned let app: FirebaseApp
3636

3737
/// A map of active `VertexAI` instances for `app`, keyed by model resource names
3838
/// (e.g., "projects/my-project-id/locations/us-central1/publishers/google/models/gemini-pro").
@@ -70,11 +70,11 @@ class VertexAIComponent: NSObject, Library, VertexAIProvider {
7070
// Unlock before the function returns.
7171
defer { os_unfair_lock_unlock(&instancesLock) }
7272

73-
if let instance = instances[app.name] {
73+
if let instance = instances[region] {
7474
return instance
7575
}
7676
let newInstance = VertexAI(app: app, region: region)
77-
instances[app.name] = newInstance
77+
instances[region] = newInstance
7878
return newInstance
7979
}
8080
}
Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
// Copyright 2024 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+
15+
import Foundation
16+
17+
import FirebaseCore
18+
@testable import FirebaseVertexAI
19+
20+
import SharedTestUtilities
21+
22+
import XCTest
23+
24+
@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
25+
class VertexComponentTests: XCTestCase {
26+
static var app: FirebaseApp!
27+
28+
override class func setUp() {
29+
super.setUp()
30+
if app == nil {
31+
let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
32+
gcmSenderID: "00000000000000000-00000000000-000000000")
33+
options.projectID = "myProjectID"
34+
FirebaseApp.configure(options: options)
35+
app = FirebaseApp(instanceWithName: "test", options: options)
36+
}
37+
}
38+
39+
// MARK: Interoperability Tests
40+
41+
/// Tests that the right number of components are being provided for the container.
42+
func testComponentsBeingRegistered() throws {
43+
let components = VertexAIComponent.componentsToRegister()
44+
XCTAssert(components.count == 1)
45+
}
46+
47+
/// Tests that a vertex instance can be created properly by the VertexAIComponent.
48+
func testVertexInstanceCreation() throws {
49+
let app = try XCTUnwrap(VertexComponentTests.app)
50+
let component = VertexAIComponent(app: app)
51+
let vertex = component.vertexAI("my-region")
52+
XCTAssertNotNil(vertex)
53+
}
54+
55+
/// Tests that the component container caches instances of VertexAIComponent.
56+
func testMultipleComponentInstancesCreated() throws {
57+
let registrants = NSMutableSet(array: [VertexAIComponent.self])
58+
let container = FirebaseComponentContainer(
59+
app: VertexComponentTests.app,
60+
registrants: registrants
61+
)
62+
63+
let provider1 = ComponentType<VertexAIProvider>.instance(for: VertexAIProvider.self,
64+
in: container)
65+
XCTAssertNotNil(provider1)
66+
67+
let provider2 = ComponentType<VertexAIProvider>.instance(for: VertexAIProvider.self,
68+
in: container)
69+
XCTAssertNotNil(provider2)
70+
71+
// Ensure they're the same instance.
72+
XCTAssert(provider1 === provider2)
73+
}
74+
75+
/// Tests that instances of vertex created are different.
76+
func testMultipleVertexInstancesCreated() throws {
77+
let app = try XCTUnwrap(VertexComponentTests.app)
78+
let registrants = NSMutableSet(array: [VertexAIComponent.self])
79+
let container = FirebaseComponentContainer(app: app, registrants: registrants)
80+
81+
let provider = ComponentType<VertexAIProvider>.instance(for: VertexAIProvider.self,
82+
in: container)
83+
XCTAssertNotNil(provider)
84+
85+
let vertex1 = provider?.vertexAI("randomRegion")
86+
let vertex2 = provider?.vertexAI("randomRegion")
87+
XCTAssertNotNil(vertex1)
88+
89+
// Ensure they're the same instance.
90+
XCTAssert(vertex1 === vertex2)
91+
92+
let vertex3 = provider?.vertexAI("differentRegion")
93+
XCTAssertNotNil(vertex3)
94+
95+
XCTAssert(vertex1 !== vertex3)
96+
}
97+
98+
/// Test that vertex instances get deallocated.
99+
func testVertexLifecycle() throws {
100+
weak var weakApp: FirebaseApp?
101+
weak var weakVertex: VertexAI?
102+
try autoreleasepool {
103+
let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
104+
gcmSenderID: "00000000000000000-00000000000-000000000")
105+
options.projectID = "myProjectID"
106+
let app1 = FirebaseApp(instanceWithName: "transitory app", options: options)
107+
weakApp = try XCTUnwrap(app1)
108+
let vertex = VertexAI(app: app1, region: "transitory bucket")
109+
weakVertex = vertex
110+
XCTAssertNotNil(weakVertex)
111+
}
112+
XCTAssertNil(weakApp)
113+
XCTAssertNil(weakVertex)
114+
}
115+
}

Package.swift

Lines changed: 12 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -170,10 +170,16 @@ let package = Package(
170170
),
171171
abseilDependency(),
172172
grpcDependency(),
173+
// TODO: restore OCMock when https://github.com/erikdoe/ocmock/pull/537
174+
// gets merged to fix Xcode 15.3 builds.
173175
.package(
174-
url: "https://github.com/erikdoe/ocmock.git",
175-
revision: "c5eeaa6dde7c308a5ce48ae4d4530462dd3a1110"
176+
url: "https://github.com/paulb777/ocmock.git",
177+
revision: "173955e93e6ee6999a10729ab67e4b4efdd1db6d"
176178
),
179+
// .package(
180+
// url: "https://github.com/erikdoe/ocmock.git",
181+
// revision: "c5eeaa6dde7c308a5ce48ae4d4530462dd3a1110"
182+
// ),
177183
.package(
178184
url: "https://github.com/firebase/leveldb.git",
179185
"1.22.2" ..< "1.23.0"
@@ -1365,11 +1371,14 @@ let package = Package(
13651371
),
13661372
.testTarget(
13671373
name: "FirebaseVertexAIUnit",
1368-
dependencies: ["FirebaseVertexAI"],
1374+
dependencies: ["FirebaseVertexAI", "SharedTestUtilities"],
13691375
path: "FirebaseVertexAI/Tests/Unit",
13701376
resources: [
13711377
.process("CountTokenResponses"),
13721378
.process("GenerateContentResponses"),
1379+
],
1380+
cSettings: [
1381+
.headerSearchPath("../../../"),
13731382
]
13741383
),
13751384
] + firestoreTargets(),

0 commit comments

Comments
 (0)