Increase Firebase AI Logic unit test coverage

FirebaseAI/Tests/Unit/ChatTests.swift

@@ -94,4 +94,103 @@ 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)
+    XCTAssertFalse(response.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)
+  }
 }

FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests+Coverage.swift

@@ -0,0 +1,124 @@
+// 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, *)
+extension GenerativeModelGoogleAITests {
+  // MARK: - Tool/Function Calling
+
+  func testGenerateContent_success_functionCall_emptyArguments() async throws {
+    MockURLProtocol
+      .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+        forResource: "unary-success-function-call-empty-arguments",
+        withExtension: "json",
+        subdirectory: "mock-responses/vertexai"
+      )
+
+    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 part = try XCTUnwrap(candidate.content.parts.first)
+    guard let functionCall = part as? FunctionCallPart else {
+      XCTFail("Part is not a FunctionCall.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, "current_time")
+    XCTAssertTrue(functionCall.args.isEmpty)
+    XCTAssertEqual(response.functionCalls, [functionCall])
+  }
+
+  func testGenerateContent_success_functionCall_withArguments() async throws {
+    MockURLProtocol
+      .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+        forResource: "unary-success-function-call-with-arguments",
+        withExtension: "json",
+        subdirectory: "mock-responses/vertexai"
+      )
+
+    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 part = try XCTUnwrap(candidate.content.parts.first)
+    guard let functionCall = part as? FunctionCallPart else {
+      XCTFail("Part is not a FunctionCall.")
+      return
+    }
+    XCTAssertEqual(functionCall.name, "sum")
+    XCTAssertEqual(functionCall.args.count, 2)
+    let argX = try XCTUnwrap(functionCall.args["x"])
+    XCTAssertEqual(argX, .number(4))
+    let argY = try XCTUnwrap(functionCall.args["y"])
+    XCTAssertEqual(argY, .number(5))
+    XCTAssertEqual(response.functionCalls, [functionCall])
+  }
+
+  func testGenerateContent_success_functionCall_parallelCalls() async throws {
+    MockURLProtocol
+      .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+        forResource: "unary-success-function-call-parallel-calls",
+        withExtension: "json",
+        subdirectory: "mock-responses/vertexai"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 3)
+    let functionCalls = response.functionCalls
+    XCTAssertEqual(functionCalls.count, 3)
+  }
+
+  func testGenerateContent_success_functionCall_mixedContent() async throws {
+    MockURLProtocol
+      .requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+        forResource: "unary-success-function-call-mixed-content",
+        withExtension: "json",
+        subdirectory: "mock-responses/vertexai"
+      )
+
+    let response = try await model.generateContent(testPrompt)
+
+    XCTAssertEqual(response.candidates.count, 1)
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertEqual(candidate.content.parts.count, 4)
+    let functionCalls = response.functionCalls
+    XCTAssertEqual(functionCalls.count, 2)
+    let text = try XCTUnwrap(response.text)
+    XCTAssertEqual(text, "The sum of [1, 2, 3] is")
+  }
+
+  // MARK: - Count Tokens
+
+  func testCountTokens_success() async throws {
+    MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
+      forResource: "unary-success-total-tokens",
+      withExtension: "json",
+      subdirectory: "mock-responses/vertexai"
+    )
+
+    let response = try await model.countTokens("Why is the sky blue?")
+    XCTAssertEqual(response.totalTokens, 6)
+  }
+}

FirebaseAI/Tests/Unit/JSONValueTests+Coverage.swift

@@ -0,0 +1,96 @@
+// 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, *)
+extension JSONValueTests {
+  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 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 json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(
+      json,
+      "{\"numberKey\":\(numberValueEncoded),\"objectKey\":{\"nestedKey\":\"nestedValue\"}}"
+    )
+  }
+
+  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 json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
+    XCTAssertEqual(
+      json,
+      "{\"arrayKey\":[\"a\",\"b\"],\"numberKey\":\(numberValueEncoded)}"
+    )
+  }
+}