empty urlMetadata array tests

Author: paulb777

FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift

@@ -380,10 +380,31 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     let successURLMetadata = urlContextMetadata.urlMetadata[1]
     XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
 
-    let errorURLMetadata = urlContextMetadata.urlMetadata[2]
+    let errorURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[2])
     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

@@ -555,6 +555,27 @@ 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

@@ -21,6 +21,12 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
     URLResponse,
     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.")
@@ -33,34 +39,44 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
   override class func canonicalRequest(for request: URLRequest) -> URLRequest { return request }
 
   override func startLoading() {
-    guard let requestHandler = MockURLProtocol.requestHandler else {
-      fatalError("`requestHandler` is nil.")
-    }
     guard let client = client else {
       fatalError("`client` is nil.")
     }
 
-    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.")
+    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)!)
             }
-            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)")
           }
-        } 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)
       }
-      client.urlProtocolDidFinishLoading(self)
+    } else {
+      fatalError("No request handler set.")
     }
   }
 

FirebaseAI/Tests/Unit/TestUtilities/GenerativeModelTestUtil.swift

@@ -79,6 +79,31 @@ enum GenerativeModelTestUtil {
     #endif // os(watchOS)
   }
 
+  /// Returns an HTTP request handler that returns the given `json` string as a response.
+  static func httpRequestHandler(json: String,
+                                 statusCode: Int = 200) throws -> ((URLRequest) throws -> (
+    URLResponse,
+    Data?
+  )) {
+    // Skip tests using MockURLProtocol on watchOS; unsupported in watchOS 2 and later, see
+    // https://developer.apple.com/documentation/foundation/urlprotocol for details.
+    #if os(watchOS)
+      throw XCTSkip("Custom URL protocols are unsupported in watchOS 2 and later.")
+    #else // os(watchOS)
+      let data = try XCTUnwrap(json.data(using: .utf8))
+      return { request in
+        let requestURL = try XCTUnwrap(request.url)
+        let response = try XCTUnwrap(HTTPURLResponse(
+          url: requestURL,
+          statusCode: statusCode,
+          httpVersion: nil,
+          headerFields: nil
+        ))
+        return (response, data)
+      }
+    #endif // os(watchOS)
+  }
+
   static func nonHTTPRequestHandler() throws -> ((URLRequest) -> (
     URLResponse,
     AsyncLineSequence<URL.AsyncBytes>?