mirror of
https://github.com/firebase/firebase-ios-sdk
synced 2026-11-11 22:55:59 +05:45
393 lines
15 KiB
Swift
393 lines
15 KiB
Swift
// Copyright 2025 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 FirebaseAppCheckInterop
|
|
import FirebaseAuthInterop
|
|
import FirebaseCore
|
|
import XCTest
|
|
|
|
@testable import FirebaseAI
|
|
|
|
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
|
|
final class GenerativeModelGoogleAITests: XCTestCase {
|
|
let testPrompt = "What sorts of questions can I ask you?"
|
|
let safetyRatingsNegligible: [SafetyRating] = [
|
|
.init(
|
|
category: .sexuallyExplicit,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .hateSpeech,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .harassment,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
.init(
|
|
category: .dangerousContent,
|
|
probability: .negligible,
|
|
probabilityScore: 0.0,
|
|
severity: SafetyRating.HarmSeverity(rawValue: "HARM_SEVERITY_UNSPECIFIED"),
|
|
severityScore: 0.0,
|
|
blocked: false
|
|
),
|
|
].sorted()
|
|
let testModelName = "test-model"
|
|
let testModelResourceName = "projects/test-project-id/models/test-model"
|
|
let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
|
|
|
|
let googleAISubdirectory = "mock-responses/googleai"
|
|
|
|
var urlSession: URLSession!
|
|
var model: GenerativeModel!
|
|
|
|
override func setUp() async throws {
|
|
let configuration = URLSessionConfiguration.default
|
|
configuration.protocolClasses = [MockURLProtocol.self]
|
|
urlSession = try XCTUnwrap(URLSession(configuration: configuration))
|
|
model = GenerativeModel(
|
|
modelName: testModelName,
|
|
modelResourceName: testModelResourceName,
|
|
firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(),
|
|
apiConfig: apiConfig,
|
|
tools: nil,
|
|
requestOptions: RequestOptions(),
|
|
urlSession: urlSession
|
|
)
|
|
}
|
|
|
|
override func tearDown() {
|
|
MockURLProtocol.requestHandler = nil
|
|
}
|
|
|
|
// MARK: - Generate Content
|
|
|
|
func testGenerateContent_success_basicReplyLong() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-long",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let finishReason = try XCTUnwrap(candidate.finishReason)
|
|
XCTAssertEqual(finishReason, .stop)
|
|
XCTAssertEqual(candidate.safetyRatings.count, 4)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let part = try XCTUnwrap(candidate.content.parts.first)
|
|
let partText = try XCTUnwrap(part as? TextPart).text
|
|
XCTAssertTrue(partText.hasPrefix("Making professional-quality"))
|
|
XCTAssertEqual(response.text, partText)
|
|
XCTAssertEqual(response.functionCalls, [])
|
|
XCTAssertEqual(response.inlineDataParts, [])
|
|
}
|
|
|
|
func testGenerateContent_success_basicReplyShort() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-short",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
let finishReason = try XCTUnwrap(candidate.finishReason)
|
|
XCTAssertEqual(finishReason, .stop)
|
|
XCTAssertEqual(candidate.safetyRatings.sorted(), safetyRatingsNegligible)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let part = try XCTUnwrap(candidate.content.parts.first)
|
|
let textPart = try XCTUnwrap(part as? TextPart)
|
|
XCTAssertTrue(textPart.text.hasPrefix("Google's headquarters"))
|
|
XCTAssertEqual(response.text, textPart.text)
|
|
XCTAssertEqual(response.functionCalls, [])
|
|
XCTAssertEqual(response.inlineDataParts, [])
|
|
}
|
|
|
|
func testGenerateContent_success_citations() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-citations",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
XCTAssertEqual(response.candidates.count, 1)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.content.parts.count, 1)
|
|
let text = try XCTUnwrap(response.text)
|
|
XCTAssertTrue(text.hasPrefix("Okay, let's break down quantum mechanics."))
|
|
let citationMetadata = try XCTUnwrap(candidate.citationMetadata)
|
|
XCTAssertEqual(citationMetadata.citations.count, 4)
|
|
let citationSource1 = try XCTUnwrap(citationMetadata.citations[0])
|
|
XCTAssertEqual(citationSource1.uri, "https://www.example.com/some-citation-1")
|
|
XCTAssertEqual(citationSource1.startIndex, 548)
|
|
XCTAssertEqual(citationSource1.endIndex, 690)
|
|
XCTAssertNil(citationSource1.title)
|
|
XCTAssertEqual(citationSource1.license, "mit")
|
|
XCTAssertNil(citationSource1.publicationDate)
|
|
let citationSource2 = try XCTUnwrap(citationMetadata.citations[1])
|
|
XCTAssertEqual(citationSource2.uri, "https://www.example.com/some-citation-1")
|
|
XCTAssertEqual(citationSource2.startIndex, 1240)
|
|
XCTAssertEqual(citationSource2.endIndex, 1407)
|
|
XCTAssertNil(citationSource2.title, "some-citation-2")
|
|
XCTAssertNil(citationSource2.license)
|
|
XCTAssertNil(citationSource2.publicationDate)
|
|
let citationSource3 = try XCTUnwrap(citationMetadata.citations[2])
|
|
XCTAssertEqual(citationSource3.startIndex, 1942)
|
|
XCTAssertEqual(citationSource3.endIndex, 2149)
|
|
XCTAssertNil(citationSource3.uri)
|
|
XCTAssertNil(citationSource3.license)
|
|
XCTAssertNil(citationSource3.title)
|
|
XCTAssertNil(citationSource3.publicationDate)
|
|
let citationSource4 = try XCTUnwrap(citationMetadata.citations[3])
|
|
XCTAssertEqual(citationSource4.startIndex, 2036)
|
|
XCTAssertEqual(citationSource4.endIndex, 2175)
|
|
XCTAssertNil(citationSource4.uri)
|
|
XCTAssertNil(citationSource4.license)
|
|
XCTAssertNil(citationSource4.title)
|
|
XCTAssertNil(citationSource4.publicationDate)
|
|
}
|
|
|
|
func testGenerateContent_usageMetadata() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-success-basic-reply-short",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let response = try await model.generateContent(testPrompt)
|
|
|
|
let usageMetadata = try XCTUnwrap(response.usageMetadata)
|
|
XCTAssertEqual(usageMetadata.promptTokenCount, 7)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails[0].modality, .text)
|
|
XCTAssertEqual(usageMetadata.promptTokensDetails[0].tokenCount, 7)
|
|
XCTAssertEqual(usageMetadata.candidatesTokenCount, 22)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails.count, 1)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].modality, .text)
|
|
XCTAssertEqual(usageMetadata.candidatesTokensDetails[0].tokenCount, 22)
|
|
}
|
|
|
|
func testGenerateContent_failure_invalidAPIKey() async throws {
|
|
let expectedStatusCode = 400
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-api-key",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory,
|
|
statusCode: expectedStatusCode
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
|
|
} catch let GenerateContentError.internalError(error as BackendError) {
|
|
XCTAssertEqual(error.httpResponseCode, 400)
|
|
XCTAssertEqual(error.status, .invalidArgument)
|
|
XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
|
|
XCTAssertTrue(error.localizedDescription.contains(error.message))
|
|
XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
|
|
XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
|
|
let nsError = error as NSError
|
|
XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
|
|
XCTAssertEqual(nsError.code, error.httpResponseCode)
|
|
return
|
|
} catch {
|
|
XCTFail("Should throw GenerateContentError.internalError(RPCError); error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
func testGenerateContent_failure_finishReasonSafety() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-finish-reason-safety",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw")
|
|
} catch let GenerateContentError.responseStoppedEarly(reason, response) {
|
|
XCTAssertEqual(reason, .safety)
|
|
XCTAssertEqual(response.text, "Safety error incoming in 5, 4, 3, 2...")
|
|
} catch {
|
|
XCTFail("Should throw a responseStoppedEarly")
|
|
}
|
|
}
|
|
|
|
func testGenerateContent_failure_unknownModel() async throws {
|
|
let expectedStatusCode = 404
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-unknown-model",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory,
|
|
statusCode: 404
|
|
)
|
|
|
|
do {
|
|
_ = try await model.generateContent(testPrompt)
|
|
XCTFail("Should throw GenerateContentError.internalError; no error thrown.")
|
|
} catch let GenerateContentError.internalError(underlying: rpcError as BackendError) {
|
|
XCTAssertEqual(rpcError.status, .notFound)
|
|
XCTAssertEqual(rpcError.httpResponseCode, expectedStatusCode)
|
|
XCTAssertTrue(rpcError.message.hasPrefix("models/gemini-5.0-flash is not found"))
|
|
} catch {
|
|
XCTFail("Should throw GenerateContentError.internalError; error thrown: \(error)")
|
|
}
|
|
}
|
|
|
|
// MARK: - Generate Content (Streaming)
|
|
|
|
func testGenerateContentStream_successBasicReplyLong() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-basic-reply-long",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = 0
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await content in stream {
|
|
XCTAssertNotNil(content.text)
|
|
responses += 1
|
|
}
|
|
|
|
XCTAssertEqual(responses, 36)
|
|
}
|
|
|
|
func testGenerateContentStream_successBasicReplyShort() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-basic-reply-short",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = 0
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await content in stream {
|
|
XCTAssertNotNil(content.text)
|
|
responses += 1
|
|
}
|
|
|
|
XCTAssertEqual(responses, 3)
|
|
}
|
|
|
|
func testGenerateContentStream_successWithCitations() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-success-citations",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
let stream = try model.generateContentStream("Hi")
|
|
var citations = [Citation]()
|
|
var responses = [GenerateContentResponse]()
|
|
for try await content in stream {
|
|
responses.append(content)
|
|
XCTAssertNotNil(content.text)
|
|
let candidate = try XCTUnwrap(content.candidates.first)
|
|
if let sources = candidate.citationMetadata?.citations {
|
|
citations.append(contentsOf: sources)
|
|
}
|
|
}
|
|
|
|
let lastCandidate = try XCTUnwrap(responses.last?.candidates.first)
|
|
XCTAssertEqual(lastCandidate.finishReason, .stop)
|
|
XCTAssertEqual(citations.count, 1)
|
|
let citation = try XCTUnwrap(citations.first)
|
|
XCTAssertEqual(citation.startIndex, 111)
|
|
XCTAssertEqual(citation.endIndex, 236)
|
|
let citationURI = try XCTUnwrap(citation.uri)
|
|
XCTAssertTrue(citationURI.starts(with: "https://www."))
|
|
XCTAssertNil(citation.license)
|
|
XCTAssertNil(citation.title)
|
|
XCTAssertNil(citation.publicationDate)
|
|
}
|
|
|
|
func testGenerateContentStream_failureInvalidAPIKey() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "unary-failure-api-key",
|
|
withExtension: "json",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
do {
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await _ in stream {
|
|
XCTFail("No content is there, this shouldn't happen.")
|
|
}
|
|
} catch let GenerateContentError.internalError(error as BackendError) {
|
|
XCTAssertEqual(error.httpResponseCode, 400)
|
|
XCTAssertEqual(error.status, .invalidArgument)
|
|
XCTAssertEqual(error.message, "API key not valid. Please pass a valid API key.")
|
|
XCTAssertTrue(error.localizedDescription.contains(error.message))
|
|
XCTAssertTrue(error.localizedDescription.contains(error.status.rawValue))
|
|
XCTAssertTrue(error.localizedDescription.contains("\(error.httpResponseCode)"))
|
|
let nsError = error as NSError
|
|
XCTAssertEqual(nsError.domain, "\(Constants.baseErrorDomain).\(BackendError.self)")
|
|
XCTAssertEqual(nsError.code, error.httpResponseCode)
|
|
return
|
|
}
|
|
|
|
XCTFail("Should have caught an error.")
|
|
}
|
|
|
|
func testGenerateContentStream_failureFinishRecitation() async throws {
|
|
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
|
|
forResource: "streaming-failure-recitation-no-content",
|
|
withExtension: "txt",
|
|
subdirectory: googleAISubdirectory
|
|
)
|
|
|
|
var responses = [GenerateContentResponse]()
|
|
do {
|
|
let stream = try model.generateContentStream("Hi")
|
|
for try await response in stream {
|
|
responses.append(response)
|
|
}
|
|
XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got no error.")
|
|
} catch let GenerateContentError.responseStoppedEarly(reason, response) {
|
|
XCTAssertEqual(reason, .recitation)
|
|
let candidate = try XCTUnwrap(response.candidates.first)
|
|
XCTAssertEqual(candidate.finishReason, reason)
|
|
} catch {
|
|
XCTFail("Expected a GenerateContentError.responseStoppedEarly error, but got error: \(error)")
|
|
}
|
|
|
|
XCTAssertEqual(responses.count, 8)
|
|
let firstResponse = try XCTUnwrap(responses.first)
|
|
XCTAssertEqual(firstResponse.text, "text1")
|
|
let lastResponse = try XCTUnwrap(responses.last)
|
|
XCTAssertEqual(lastResponse.text, "text8")
|
|
}
|
|
}
|