From d9d5fa7ca00c45526babae63685b9a413b797d0a Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 5 Jun 2024 16:01:54 -0400 Subject: [PATCH] [Vertex AI] Add simple integration tests --- .../Tests/Integration/IntegrationTests.swift | 68 ++++++++++++++++ .../Integration/Resources/placeholder.txt | 0 Package.swift | 8 ++ .../vertexai-integration.plist.gpg | Bin 0 -> 619 bytes .../FirebaseVertexAIIntegration.xcscheme | 77 ++++++++++++++++++ 5 files changed, 153 insertions(+) create mode 100644 FirebaseVertexAI/Tests/Integration/IntegrationTests.swift create mode 100644 FirebaseVertexAI/Tests/Integration/Resources/placeholder.txt create mode 100644 scripts/gha-encrypted/vertexai-integration.plist.gpg create mode 100644 scripts/spm_test_schemes/FirebaseVertexAIIntegration.xcscheme 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 0000000000000000000000000000000000000000..e37af69add52c19ba4c4dd6a02831d7b1e44efa5 GIT binary patch literal 619 zcmV-x0+juX4Fm}T2=Qx<1+%K6^!w830g%u+trCsBD~NCo#p#~Nc4E7mTbfC9U>^{E zMO&#K!U-YKe!dmuH0@x%8uH?;=$dPj3xH2PL15WEJ zKK=e;Yyd!gzy!HFI0P8(HDE2bOxL!rmYi^H2v182nj_u%RWJsUEkCdpSVn4TrI0n` zh8Fui!SnU_*7@`lb}hq=)qSQ>rVyx7vT%j!pm`hHjIOlpf*tLFZBguZr{tso74%ZB zA2@oAW%E@#dAe>2ZYRtihXyTbxYa=u9TLmnikS~dMQDpLLFVNYot}&w*@n#V7OHnY zK)n{GRd2B_NjNVQHBU3X0ybP~RrCD7$G0)SKnMMZDk0ij;OKLg^7LGP*p@7m{1|JG z`MRmIFJq1iHg=U`Gz4ax4pk7^snz8Ra#Pqk6@`C#XUIVLGSY4~$#Z%er&}4SR(K`$ zNcn_AyAqm~P57!DXn?_*9pc#1k6i$5Idet)9C*XtX>rE0H%+h`$*OKQ$0ME0$&Oyft>MTD;nV0 zA9BFMfC=G#)Wa;_1%Z3XS#VXgN&X6@tZPobaTHxoZS2G=^lpj8yeh5c)!_VQ)oVOs`XhNIw0{pLO?YCNJigrk=gxvH>O339Hz#30TnU9!?Tw|C3II~ZIhz=%Yx F_??J|FiQXc literal 0 HcmV?d00001 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 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + + + + +