diff --git a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift new file mode 100644 index 00000000000..2ff98ee0e49 --- /dev/null +++ b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift @@ -0,0 +1,68 @@ +// Copyright 2024 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import FirebaseCore +import FirebaseVertexAI +import XCTest + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, *) +final class IntegrationTests: XCTestCase { + let generationConfig = GenerationConfig(temperature: 0.0, topP: 0.0, topK: 1) + + var vertex: VertexAI! + var model: GenerativeModel! + + override func setUp() async throws { + let plistPath = try XCTUnwrap(Bundle.module.path( + forResource: "GoogleService-Info", + ofType: "plist" + )) + let options = try XCTUnwrap(FirebaseOptions(contentsOfFile: plistPath)) + FirebaseApp.configure(options: options) + + vertex = VertexAI.vertexAI() + model = vertex.generativeModel( + modelName: "gemini-1.5-flash", + generationConfig: generationConfig + ) + } + + override func tearDown() async throws { + if let app = FirebaseApp.app() { + await app.delete() + } + } + + // MARK: - Generate Content + + func testGenerateContent() async throws { + let prompt = "Where is Google headquarters located? Answer with the city name only." + + let response = try await model.generateContent(prompt) + + let text = try XCTUnwrap(response.text).trimmingCharacters(in: .whitespacesAndNewlines) + XCTAssertEqual(text, "Mountain View") + } + + // MARK: - Count Tokens + + func testCountTokens() async throws { + let prompt = "Why is the sky blue?" + + let response = try await model.countTokens(prompt) + + XCTAssertEqual(response.totalTokens, 6) + XCTAssertEqual(response.totalBillableCharacters, 16) + } +} diff --git a/FirebaseVertexAI/Tests/Integration/Resources/placeholder.txt b/FirebaseVertexAI/Tests/Integration/Resources/placeholder.txt new file mode 100644 index 00000000000..e69de29bb2d diff --git a/Package.swift b/Package.swift index 5d58d0637f8..f104e41e747 100644 --- a/Package.swift +++ b/Package.swift @@ -1386,6 +1386,14 @@ let package = Package( .headerSearchPath("../../../"), ] ), + .testTarget( + name: "FirebaseVertexAIIntegration", + dependencies: ["FirebaseVertexAI", "SharedTestUtilities"], + path: "FirebaseVertexAI/Tests/Integration", + resources: [ + .process("Resources"), + ] + ), ] + firestoreTargets(), cLanguageStandard: .c99, cxxLanguageStandard: CXXLanguageStandard.gnucxx14 diff --git a/scripts/gha-encrypted/vertexai-integration.plist.gpg b/scripts/gha-encrypted/vertexai-integration.plist.gpg new file mode 100644 index 00000000000..e37af69add5 Binary files /dev/null and b/scripts/gha-encrypted/vertexai-integration.plist.gpg differ diff --git a/scripts/spm_test_schemes/FirebaseVertexAIIntegration.xcscheme b/scripts/spm_test_schemes/FirebaseVertexAIIntegration.xcscheme new file mode 100644 index 00000000000..b2658440824 --- /dev/null +++ b/scripts/spm_test_schemes/FirebaseVertexAIIntegration.xcscheme @@ -0,0 +1,77 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +