update Mock to handle errors

fumito-ito · Oct 24, 2024 · f09c437 · f09c437
1 parent fcf5db3
commit f09c437
Show file tree

Hide file tree

Showing 2 changed files with 36 additions and 3 deletions.
diff --git a/Sources/AnthropicSwiftSDK-TestUtils/HTTPMock.swift b/Sources/AnthropicSwiftSDK-TestUtils/HTTPMock.swift
@@ -12,6 +12,7 @@ public enum MockInspectType {
     case request((URLRequest) -> Void, String?)
     case requestHeader(([String: String]?) -> Void, String?)
     case response(String)
+    case error(String)
 }
 
 public class HTTPMock: URLProtocol {
@@ -35,24 +36,37 @@ public class HTTPMock: URLProtocol {
         }
 
         if let url = request.url,
-           let response = HTTPURLResponse(url: url, statusCode: 200, httpVersion: "HTTP/2", headerFields: nil) {
-            client?.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)
+           let succeedResponse = HTTPURLResponse(url: url, statusCode: 200, httpVersion: "HTTP/2", headerFields: nil),
+           let errorResponse = HTTPURLResponse(url: url, statusCode: 400, httpVersion: "HTTP/2", headerFields: nil) {
 
             switch Self.inspectType {
             case .none:
+                client?.urlProtocol(self, didReceive: succeedResponse, cacheStoragePolicy: .notAllowed)
                 client?.urlProtocol(self, didLoad: getBasicJSONStringData())
+            case .request(_, nil), .requestHeader(_, nil):
+                client?.urlProtocol(self, didReceive: succeedResponse, cacheStoragePolicy: .notAllowed)
             case .request(_, let jsonString), .requestHeader(_, let jsonString):
+                client?.urlProtocol(self, didReceive: succeedResponse, cacheStoragePolicy: .notAllowed)
                 guard let jsonString, let data = jsonString.data(using: .utf8) else {
                     client?.urlProtocol(self, didLoad: getBasicJSONStringData())
                     return
                 }
                 client?.urlProtocol(self, didLoad: data)
             case .response(let jsonString):
+                client?.urlProtocol(self, didReceive: succeedResponse, cacheStoragePolicy: .notAllowed)
                 guard let data = jsonString.data(using: .utf8) else {
                     client?.urlProtocol(self, didLoad: getBasicJSONStringData())
                     return
                 }
                 client?.urlProtocol(self, didLoad: data)
+            case .error(let responseJSON):
+                client?.urlProtocol(self, didReceive: errorResponse, cacheStoragePolicy: .notAllowed)
+                guard
+                    let data = responseJSON.data(using: .utf8) else {
+                    client?.urlProtocol(self, didLoad: getBasicJSONStringData())
+                    return
+                }
+                client?.urlProtocol(self, didLoad: data)
             }
         }
 

diff --git a/Tests/AnthropicSwiftSDKTests/MessagesTests.swift b/Tests/AnthropicSwiftSDKTests/MessagesTests.swift
@@ -26,7 +26,26 @@ final class MessagesTests: XCTestCase {
         HTTPMock.inspectType = .requestHeader({ headers in
             XCTAssertEqual(headers!["x-api-key"], "This-is-test-API-key")
             expectation.fulfill()
-        }, nil)
+        }, """
+        {
+          "id": "msg_01XFDUDYJgAACzvnptvVoYEL",
+          "type": "message",
+          "role": "assistant",
+          "content": [
+            {
+              "type": "text",
+              "text": "Hello!"
+            }
+          ],
+          "model": "claude-2.1",
+          "stop_reason": "end_turn",
+          "stop_sequence": null,
+          "usage": {
+            "input_tokens": 12,
+            "output_tokens": 6
+          }
+        }
+        """)
 
         let message = Message(role: .user, content: [.text("This is test text")])
         let _ = try await messages.createMessage([message], maxTokens: 1024)