diff --git a/FirebaseAI/Tests/Unit/ChatTests.swift b/FirebaseAI/Tests/Unit/ChatTests.swift index 40373a47494..4bf89f8cb51 100644 --- a/FirebaseAI/Tests/Unit/ChatTests.swift +++ b/FirebaseAI/Tests/Unit/ChatTests.swift @@ -94,4 +94,104 @@ final class ChatTests: XCTestCase { XCTAssertEqual(chat.history[1], assembledExpectation) #endif // os(watchOS) } + + func testSendMessage_unary_appendsHistory() async throws { + let expectedInput = "Test input" + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "unary-success-basic-reply-short", + withExtension: "json", + subdirectory: "mock-responses/googleai" + ) + let model = GenerativeModel( + modelName: modelName, + modelResourceName: modelResourceName, + firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(), + apiConfig: FirebaseAI.defaultVertexAIAPIConfig, + tools: nil, + requestOptions: RequestOptions(), + urlSession: urlSession + ) + let chat = model.startChat() + + // Pre-condition: History should be empty. + XCTAssertTrue(chat.history.isEmpty) + + let response = try await chat.sendMessage(expectedInput) + + XCTAssertNotNil(response.text) + let text = try XCTUnwrap(response.text) + XCTAssertFalse(text.isEmpty) + + // Post-condition: History should have the user's message and the model's response. + XCTAssertEqual(chat.history.count, 2) + let userInput = try XCTUnwrap(chat.history.first) + XCTAssertEqual(userInput.role, "user") + XCTAssertEqual(userInput.parts.count, 1) + let userInputText = try XCTUnwrap(userInput.parts.first as? TextPart) + XCTAssertEqual(userInputText.text, expectedInput) + + let modelResponse = try XCTUnwrap(chat.history.last) + XCTAssertEqual(modelResponse.role, "model") + XCTAssertEqual(modelResponse.parts.count, 1) + let modelResponseText = try XCTUnwrap(modelResponse.parts.first as? TextPart) + XCTAssertFalse(modelResponseText.text.isEmpty) + } + + func testSendMessageStream_error_doesNotAppendHistory() async throws { + MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler( + forResource: "streaming-failure-finish-reason-safety", + withExtension: "txt", + subdirectory: "mock-responses/vertexai" + ) + let model = GenerativeModel( + modelName: modelName, + modelResourceName: modelResourceName, + firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(), + apiConfig: FirebaseAI.defaultVertexAIAPIConfig, + tools: nil, + requestOptions: RequestOptions(), + urlSession: urlSession + ) + let chat = model.startChat() + let input = "Test input" + + // Pre-condition: History should be empty. + XCTAssertTrue(chat.history.isEmpty) + + do { + let stream = try chat.sendMessageStream(input) + for try await _ in stream { + // Consume the stream. + } + XCTFail("Should have thrown a responseStoppedEarly error.") + } catch let GenerateContentError.responseStoppedEarly(reason, _) { + XCTAssertEqual(reason, .safety) + } catch { + XCTFail("Unexpected error thrown: \(error)") + } + + // Post-condition: History should still be empty. + XCTAssertEqual(chat.history.count, 0) + } + + func testStartChat_withHistory_initializesCorrectly() async throws { + let history = [ + ModelContent(role: "user", parts: "Question 1"), + ModelContent(role: "model", parts: "Answer 1"), + ] + let model = GenerativeModel( + modelName: modelName, + modelResourceName: modelResourceName, + firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(), + apiConfig: FirebaseAI.defaultVertexAIAPIConfig, + tools: nil, + requestOptions: RequestOptions(), + urlSession: urlSession + ) + + let chat = model.startChat(history: history) + + XCTAssertEqual(chat.history.count, 2) + XCTAssertEqual(chat.history, history) + } } diff --git a/FirebaseAI/Tests/Unit/JSONValueTests.swift b/FirebaseAI/Tests/Unit/JSONValueTests.swift index 1ffe88eaf55..54ac3520e77 100644 --- a/FirebaseAI/Tests/Unit/JSONValueTests.swift +++ b/FirebaseAI/Tests/Unit/JSONValueTests.swift @@ -97,6 +97,48 @@ final class JSONValueTests: XCTestCase { XCTAssertEqual(json, "null") } + func testDecodeNestedObject() throws { + let nestedObject: JSONObject = [ + "nestedKey": .string("nestedValue"), + ] + let expectedObject: JSONObject = [ + "numberKey": .number(numberValue), + "objectKey": .object(nestedObject), + ] + let json = """ + { + "numberKey": \(numberValue), + "objectKey": { + "nestedKey": "nestedValue" + } + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let jsonObject = try XCTUnwrap(decoder.decode(JSONValue.self, from: jsonData)) + + XCTAssertEqual(jsonObject, .object(expectedObject)) + } + + func testDecodeNestedArray() throws { + let nestedArray: [JSONValue] = [.string("a"), .string("b")] + let expectedObject: JSONObject = [ + "numberKey": .number(numberValue), + "arrayKey": .array(nestedArray), + ] + let json = """ + { + "numberKey": \(numberValue), + "arrayKey": ["a", "b"] + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + + let jsonObject = try XCTUnwrap(decoder.decode(JSONValue.self, from: jsonData)) + + XCTAssertEqual(jsonObject, .object(expectedObject)) + } + func testEncodeNumber() throws { let jsonData = try encoder.encode(JSONValue.number(numberValue)) @@ -143,4 +185,30 @@ final class JSONValueTests: XCTestCase { let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) XCTAssertEqual(json, "[null,\(numberValueEncoded)]") } + + func testEncodeNestedObject() throws { + let nestedObject: JSONObject = [ + "nestedKey": .string("nestedValue"), + ] + let objectValue: JSONObject = [ + "numberKey": .number(numberValue), + "objectKey": .object(nestedObject), + ] + + let jsonData = try encoder.encode(JSONValue.object(objectValue)) + let jsonObject = try XCTUnwrap(decoder.decode(JSONValue.self, from: jsonData)) + XCTAssertEqual(jsonObject, .object(objectValue)) + } + + func testEncodeNestedArray() throws { + let nestedArray: [JSONValue] = [.string("a"), .string("b")] + let objectValue: JSONObject = [ + "numberKey": .number(numberValue), + "arrayKey": .array(nestedArray), + ] + + let jsonData = try encoder.encode(JSONValue.object(objectValue)) + let jsonObject = try XCTUnwrap(decoder.decode(JSONValue.self, from: jsonData)) + XCTAssertEqual(jsonObject, .object(objectValue)) + } } diff --git a/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift b/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift index e7531d1da9e..658db79a50e 100644 --- a/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift +++ b/FirebaseAI/Tests/Unit/PartsRepresentableTests.swift @@ -121,4 +121,55 @@ final class PartsRepresentableTests: XCTestCase { } } #endif + + func testMixedParts() throws { + let text = "This is a test" + let data = try XCTUnwrap("This is some data".data(using: .utf8)) + let inlineData = InlineDataPart(data: data, mimeType: "text/plain") + + let parts: [any PartsRepresentable] = [text, inlineData] + let modelContent = ModelContent(parts: parts) + + XCTAssertEqual(modelContent.parts.count, 2) + let textPart = try XCTUnwrap(modelContent.parts[0] as? TextPart) + XCTAssertEqual(textPart.text, text) + let dataPart = try XCTUnwrap(modelContent.parts[1] as? InlineDataPart) + XCTAssertEqual(dataPart, inlineData) + } + + #if canImport(UIKit) + func testMixedParts_withImage() throws { + let text = "This is a test" + let image = try XCTUnwrap(UIImage(systemName: "star")) + let parts: [any PartsRepresentable] = [text, image] + let modelContent = ModelContent(parts: parts) + + XCTAssertEqual(modelContent.parts.count, 2) + let textPart = try XCTUnwrap(modelContent.parts[0] as? TextPart) + XCTAssertEqual(textPart.text, text) + let imagePart = try XCTUnwrap(modelContent.parts[1] as? InlineDataPart) + XCTAssertEqual(imagePart.mimeType, "image/jpeg") + XCTAssertFalse(imagePart.data.isEmpty) + } + + #elseif canImport(AppKit) + func testMixedParts_withImage() throws { + let text = "This is a test" + let coreImage = CIImage(color: CIColor.blue) + .cropped(to: CGRect(origin: CGPoint.zero, size: CGSize(width: 16, height: 16))) + let rep = NSCIImageRep(ciImage: coreImage) + let image = NSImage(size: rep.size) + image.addRepresentation(rep) + + let parts: [any PartsRepresentable] = [text, image] + let modelContent = ModelContent(parts: parts) + + XCTAssertEqual(modelContent.parts.count, 2) + let textPart = try XCTUnwrap(modelContent.parts[0] as? TextPart) + XCTAssertEqual(textPart.text, text) + let imagePart = try XCTUnwrap(modelContent.parts[1] as? InlineDataPart) + XCTAssertEqual(imagePart.mimeType, "image/jpeg") + XCTAssertFalse(imagePart.data.isEmpty) + } + #endif } diff --git a/FirebaseAI/Tests/Unit/SafetyTests.swift b/FirebaseAI/Tests/Unit/SafetyTests.swift new file mode 100644 index 00000000000..4a1e07e04e3 --- /dev/null +++ b/FirebaseAI/Tests/Unit/SafetyTests.swift @@ -0,0 +1,123 @@ +// 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 XCTest + +@testable import FirebaseAI + +@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *) +final class SafetyTests: XCTestCase { + let decoder = JSONDecoder() + let encoder = JSONEncoder() + + override func setUp() { + encoder.outputFormatting = .init( + arrayLiteral: .prettyPrinted, .sortedKeys, .withoutEscapingSlashes + ) + } + + // MARK: - SafetyRating Decoding + + func testDecodeSafetyRating_allFieldsPresent() throws { + let json = """ + { + "category": "HARM_CATEGORY_DANGEROUS_CONTENT", + "probability": "NEGLIGIBLE", + "probabilityScore": 0.1, + "severity": "HARM_SEVERITY_LOW", + "severityScore": 0.2, + "blocked": true + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + let rating = try decoder.decode(SafetyRating.self, from: jsonData) + + XCTAssertEqual(rating.category, .dangerousContent) + XCTAssertEqual(rating.probability, .negligible) + XCTAssertEqual(rating.probabilityScore, 0.1) + XCTAssertEqual(rating.severity, .low) + XCTAssertEqual(rating.severityScore, 0.2) + XCTAssertTrue(rating.blocked) + } + + func testDecodeSafetyRating_missingOptionalFields() throws { + let json = """ + { + "category": "HARM_CATEGORY_HARASSMENT", + "probability": "LOW" + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + let rating = try decoder.decode(SafetyRating.self, from: jsonData) + + XCTAssertEqual(rating.category, .harassment) + XCTAssertEqual(rating.probability, .low) + XCTAssertEqual(rating.probabilityScore, 0.0) + XCTAssertEqual(rating.severity, .unspecified) + XCTAssertEqual(rating.severityScore, 0.0) + XCTAssertFalse(rating.blocked) + } + + func testDecodeSafetyRating_unknownEnums() throws { + let json = """ + { + "category": "HARM_CATEGORY_UNKNOWN", + "probability": "UNKNOWN_PROBABILITY", + "severity": "UNKNOWN_SEVERITY" + } + """ + let jsonData = try XCTUnwrap(json.data(using: .utf8)) + let rating = try decoder.decode(SafetyRating.self, from: jsonData) + + XCTAssertEqual(rating.category.rawValue, "HARM_CATEGORY_UNKNOWN") + XCTAssertEqual(rating.probability.rawValue, "UNKNOWN_PROBABILITY") + XCTAssertEqual(rating.severity.rawValue, "UNKNOWN_SEVERITY") + } + + // MARK: - SafetySetting Encoding + + func testEncodeSafetySetting_allFields() throws { + let setting = SafetySetting( + harmCategory: .hateSpeech, + threshold: .blockMediumAndAbove, + method: .severity + ) + let jsonData = try encoder.encode(setting) + let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + + XCTAssertEqual(jsonString, """ + { + "category" : "HARM_CATEGORY_HATE_SPEECH", + "method" : "SEVERITY", + "threshold" : "BLOCK_MEDIUM_AND_ABOVE" + } + """) + } + + func testEncodeSafetySetting_nilMethod() throws { + let setting = SafetySetting( + harmCategory: .sexuallyExplicit, + threshold: .blockOnlyHigh + ) + let jsonData = try encoder.encode(setting) + let jsonString = try XCTUnwrap(String(data: jsonData, encoding: .utf8)) + + XCTAssertEqual(jsonString, """ + { + "category" : "HARM_CATEGORY_SEXUALLY_EXPLICIT", + "threshold" : "BLOCK_ONLY_HIGH" + } + """) + } +}