Revert MockURLProtocol changes and inline JSON tests

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

@@ -106,4 +106,23 @@ final class GenerateContentResponseTests: XCTestCase {
       "functionCalls should be empty when there are no candidates."
     )
   }
+
+  func testURLContextMetadata_withEmptyURLMetadata_isNil() throws {
+    let json = """
+    {
+      "candidates": [
+        {
+          "content": { "role": "model", "parts": [ { "text": "Some text." } ] },
+          "finishReason": "STOP",
+          "urlContextMetadata": { "urlMetadata": [] }
+        }
+      ]
+    }
+    """.data(using: .utf8)!
+
+    let response = try JSONDecoder().decode(GenerateContentResponse.self, from: json)
+
+    let candidate = try XCTUnwrap(response.candidates.first)
+    XCTAssertNil(candidate.urlContextMetadata)
+  }
 }