[Vertex AI] Index instances by app name and location

firebase · Oct 31, 2024 · 0309049 · 0309049
1 parent 948c460
commit 0309049
Show file tree

Hide file tree

Showing 2 changed files with 50 additions and 22 deletions.
diff --git a/FirebaseVertexAI/Sources/VertexAI.swift b/FirebaseVertexAI/Sources/VertexAI.swift
@@ -55,11 +55,12 @@ public class VertexAI {
     // Unlock before the function returns.
     defer { os_unfair_lock_unlock(&instancesLock) }
 
-    if let instance = instances[location] {
+    let instanceKey = "\(app.name):\(location)"
+    if let instance = instances[instanceKey] {
       return instance
     }
     let newInstance = VertexAI(app: app, location: location)
-    instances[location] = newInstance
+    instances[instanceKey] = newInstance
     return newInstance
   }
 

diff --git a/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift b/FirebaseVertexAI/Tests/Unit/VertexComponentTests.swift
@@ -24,22 +24,21 @@ import XCTest
 class VertexComponentTests: XCTestCase {
   static let projectID = "test-project-id"
   static let apiKey = "test-api-key"
+  static let options = {
+    let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
+                                  gcmSenderID: "00000000000000000-00000000000-000000000")
+    options.projectID = VertexComponentTests.projectID
+    options.apiKey = VertexComponentTests.apiKey
 
-  static var app: FirebaseApp?
+    return options
+  }()
 
-  let location = "test-location"
+  static let app = {
+    FirebaseApp.configure(options: options)
+    return FirebaseApp(instanceWithName: "test", options: options)
+  }()
 
-  override class func setUp() {
-    super.setUp()
-    if app == nil {
-      let options = FirebaseOptions(googleAppID: "0:0000000000000:ios:0000000000000000",
-                                    gcmSenderID: "00000000000000000-00000000000-000000000")
-      options.projectID = VertexComponentTests.projectID
-      options.apiKey = VertexComponentTests.apiKey
-      FirebaseApp.configure(options: options)
-      app = FirebaseApp(instanceWithName: "test", options: options)
-    }
-  }
+  let location = "test-location"
 
   /// Test that the objc class is available for the component system to update the user agent.
   func testComponentsBeingRegistered() throws {
@@ -48,27 +47,55 @@ class VertexComponentTests: XCTestCase {
 
   /// Tests that a vertex instance can be created properly.
   func testVertexInstanceCreation() throws {
-    let app = try XCTUnwrap(VertexComponentTests.app)
-
-    let vertex = VertexAI.vertexAI(app: app, location: location)
+    let vertex = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
 
     XCTAssertNotNil(vertex)
     XCTAssertEqual(vertex.projectID, VertexComponentTests.projectID)
     XCTAssertEqual(vertex.apiKey, VertexComponentTests.apiKey)
     XCTAssertEqual(vertex.location, location)
   }
 
-  /// Tests that a vertex instances are reused properly.
-  func testMultipleComponentInstancesCreated() throws {
+  /// Tests that Vertex instances are reused properly.
+  func testSameAppAndLocation_instanceReused() throws {
     let app = try XCTUnwrap(VertexComponentTests.app)
+
     let vertex1 = VertexAI.vertexAI(app: app, location: location)
     let vertex2 = VertexAI.vertexAI(app: app, location: location)
 
     // Ensure they're the same instance.
     XCTAssert(vertex1 === vertex2)
+  }
+
+  func testSameAppAndDifferentLocation_newInstanceCreated() throws {
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: VertexComponentTests.app, location: "differentLocation")
+
+    // Ensure they are different instances.
+    XCTAssert(vertex1 !== vertex2)
+  }
+
+  func testDifferentAppAndSameLocation_newInstanceCreated() throws {
+    FirebaseApp.configure(name: "test-2", options: VertexComponentTests.options)
+    let app2 = FirebaseApp(instanceWithName: "test-2", options: VertexComponentTests.options)
+    addTeardownBlock { await app2.delete() }
+
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: app2, location: location)
+
+    XCTAssert(VertexComponentTests.app != app2)
+    XCTAssert(vertex1 !== vertex2) // Ensure they are different instances.
+  }
+
+  func testDifferentAppAndDifferentLocation_newInstanceCreated() throws {
+    FirebaseApp.configure(name: "test-2", options: VertexComponentTests.options)
+    let app2 = FirebaseApp(instanceWithName: "test-2", options: VertexComponentTests.options)
+    addTeardownBlock { await app2.delete() }
+
+    let vertex1 = VertexAI.vertexAI(app: VertexComponentTests.app, location: location)
+    let vertex2 = VertexAI.vertexAI(app: app2, location: "differentLocation")
 
-    let vertex3 = VertexAI.vertexAI(app: app, location: "differentLocation")
-    XCTAssert(vertex1 !== vertex3)
+    XCTAssert(VertexComponentTests.app != app2)
+    XCTAssert(vertex1 !== vertex2) // Ensure they are different instances.
   }
 
   /// Test that vertex instances get deallocated.