[Firebase AI] Add Candidate decoding tests for urlMetadata (#15348)

Author: andrewheard

FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift

@@ -384,27 +384,6 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
   }
 
-  func testGenerateContent_success_urlContext_emptyURLMetadata() async throws {
-    let json = """
-    {
-      "candidates": [
-        {
-          "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
-          "finishReason": "STOP",
-          "urlContextMetadata": { "urlMetadata": [] }
-        }
-      ]
-    }
-    """
-    MockURLProtocol.requestHandler = nil
-    MockURLProtocol.dataRequestHandler = try GenerativeModelTestUtil.httpRequestHandler(json: json)
-
-    let response = try await model.generateContent(testPrompt)
-
-    let candidate = try XCTUnwrap(response.candidates.first)
-    XCTAssertNil(candidate.urlContextMetadata)
-  }
-
   func testGenerateContent_failure_invalidAPIKey() async throws {
     let expectedStatusCode = 400
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(

FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift

@@ -556,27 +556,6 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertEqual(urlMetadata.retrievalStatus, .error)
   }
 
-  func testGenerateContent_success_urlContext_emptyURLMetadata() async throws {
-    let json = """
-    {
-      "candidates": [
-        {
-          "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
-          "finishReason": "STOP",
-          "urlContextMetadata": { "urlMetadata": [] }
-        }
-      ]
-    }
-    """
-    MockURLProtocol.requestHandler = nil
-    MockURLProtocol.dataRequestHandler = try GenerativeModelTestUtil.httpRequestHandler(json: json)
-
-    let response = try await model.generateContent(testPrompt)
-
-    let candidate = try XCTUnwrap(response.candidates.first)
-    XCTAssertNil(candidate.urlContextMetadata)
-  }
-
   func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {
     MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
       forResource: "unary-success-image-invalid-safety-ratings",

FirebaseAI/Tests/Unit/MockURLProtocol.swift

@@ -22,11 +22,6 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
     AsyncLineSequence<URL.AsyncBytes>?
   ))?
 
-  nonisolated(unsafe) static var dataRequestHandler: ((URLRequest) throws -> (
-    URLResponse,
-    Data?
-  ))?
-
   override class func canInit(with request: URLRequest) -> Bool {
     #if os(watchOS)
       print("MockURLProtocol cannot be used on watchOS.")
@@ -43,40 +38,31 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
       fatalError("`client` is nil.")
     }
 
-    if let requestHandler = MockURLProtocol.requestHandler {
-      Task {
-        let (response, stream) = try requestHandler(self.request)
-        client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
-        if let stream = stream {
-          do {
-            for try await line in stream {
-              guard let data = line.data(using: .utf8) else {
-                fatalError("Failed to convert \"\(line)\" to UTF8 data.")
-              }
-              client.urlProtocol(self, didLoad: data)
-              // Add a newline character since AsyncLineSequence strips them when reading line by
-              // line;
-              // without the following, the whole file is delivered as a single line.
-              client.urlProtocol(self, didLoad: "\n".data(using: .utf8)!)
+    guard let requestHandler = MockURLProtocol.requestHandler else {
+      fatalError("No request handler set.")
+    }
+
+    Task {
+      let (response, stream) = try requestHandler(self.request)
+      client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+      if let stream = stream {
+        do {
+          for try await line in stream {
+            guard let data = line.data(using: .utf8) else {
+              fatalError("Failed to convert \"\(line)\" to UTF8 data.")
             }
-          } catch {
-            client.urlProtocol(self, didFailWithError: error)
-            XCTFail("Unexpected failure reading lines from stream: \(error.localizedDescription)")
+            client.urlProtocol(self, didLoad: data)
+            // Add a newline character since AsyncLineSequence strips them when reading line by
+            // line;
+            // without the following, the whole file is delivered as a single line.
+            client.urlProtocol(self, didLoad: "\n".data(using: .utf8)!)
           }
+        } catch {
+          client.urlProtocol(self, didFailWithError: error)
+          XCTFail("Unexpected failure reading lines from stream: \(error.localizedDescription)")
         }
-        client.urlProtocolDidFinishLoading(self)
-      }
-    } else if let dataRequestHandler = MockURLProtocol.dataRequestHandler {
-      Task {
-        let (response, data) = try dataRequestHandler(self.request)
-        client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
-        if let data = data {
-          client.urlProtocol(self, didLoad: data)
-        }
-        client.urlProtocolDidFinishLoading(self)
       }
-    } else {
-      fatalError("No request handler set.")
+      client.urlProtocolDidFinishLoading(self)
     }
   }
 

FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift

@@ -17,6 +17,8 @@ import XCTest
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 final class GenerateContentResponseTests: XCTestCase {
+  let jsonDecoder = JSONDecoder()
+
   // MARK: - GenerateContentResponse Computed Properties
 
   func testGenerateContentResponse_inlineDataParts_success() throws {
@@ -106,4 +108,53 @@ final class GenerateContentResponseTests: XCTestCase {
       "functionCalls should be empty when there are no candidates."
     )
   }
+
+  // MARK: - Decoding Tests
+
+  func testDecodeCandidate_emptyURLMetadata_urlContextMetadataIsNil() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "STOP",
+      "urlContextMetadata": { "urlMetadata": [] }
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertNil(
+      candidate.urlContextMetadata,
+      "urlContextMetadata should be nil if the `urlMetadata` array is empty in the candidate."
+    )
+    XCTAssertEqual(candidate.content.role, "model")
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    let textPart = try XCTUnwrap(part as? TextPart)
+    XCTAssertEqual(textPart.text, "Some text.")
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(candidate.finishReason, .stop)
+  }
+
+  func testDecodeCandidate_missingURLMetadata_urlContextMetadataIsNil() throws {
+    let json = """
+    {
+      "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+      "finishReason": "STOP"
+    }
+    """
+    let jsonData = try XCTUnwrap(json.data(using: .utf8))
+
+    let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
+
+    XCTAssertNil(
+      candidate.urlContextMetadata,
+      "urlContextMetadata should be nil if `urlMetadata` is not provided in the candidate."
+    )
+    XCTAssertEqual(candidate.content.role, "model")
+    let part = try XCTUnwrap(candidate.content.parts.first)
+    let textPart = try XCTUnwrap(part as? TextPart)
+    XCTAssertEqual(textPart.text, "Some text.")
+    XCTAssertFalse(textPart.isThought)
+    XCTAssertEqual(candidate.finishReason, .stop)
+  }
 }