From f3d07ac871151da440348b72ea4afba68f5cfea3 Mon Sep 17 00:00:00 2001 From: Ryan Wilson Date: Tue, 12 Dec 2023 22:53:52 -0500 Subject: [PATCH] Further async streaming testing (#48) Includes cleanup, and moving to new naming conventions for tests. --- .../streaming-failure-error-mid-stream.txt | 17 +++ .../streaming-failure-invalid-json.txt} | 0 .../GoogleAITests/GenerativeModelTests.swift | 109 +++++------------- .../EmptyContentStreamingResponse.json | 1 - .../ExampleErrorMidStream.json | 17 --- 5 files changed, 47 insertions(+), 97 deletions(-) create mode 100644 Tests/GoogleAITests/GenerateContentResponses/streaming-failure-error-mid-stream.txt rename Tests/GoogleAITests/{SampleResponses/InvalidStreamingResponse.json => GenerateContentResponses/streaming-failure-invalid-json.txt} (100%) delete mode 100644 Tests/GoogleAITests/SampleResponses/EmptyContentStreamingResponse.json delete mode 100644 Tests/GoogleAITests/SampleResponses/ExampleErrorMidStream.json diff --git a/Tests/GoogleAITests/GenerateContentResponses/streaming-failure-error-mid-stream.txt b/Tests/GoogleAITests/GenerateContentResponses/streaming-failure-error-mid-stream.txt new file mode 100644 index 0000000..aeb4eb0 --- /dev/null +++ b/Tests/GoogleAITests/GenerateContentResponses/streaming-failure-error-mid-stream.txt @@ -0,0 +1,17 @@ +data: {"candidates": [{"content": {"parts": [{"text": "First "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} + +data: {"candidates": [{"content": {"parts": [{"text": "Second "}]},"finishReason": "STOP","index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} + +{ + "error": { + "code": 499, + "message": "The operation was cancelled.", + "status": "CANCELLED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::cancelled: " + } + ] + } +} diff --git a/Tests/GoogleAITests/SampleResponses/InvalidStreamingResponse.json b/Tests/GoogleAITests/GenerateContentResponses/streaming-failure-invalid-json.txt similarity index 100% rename from Tests/GoogleAITests/SampleResponses/InvalidStreamingResponse.json rename to Tests/GoogleAITests/GenerateContentResponses/streaming-failure-invalid-json.txt diff --git a/Tests/GoogleAITests/GenerativeModelTests.swift b/Tests/GoogleAITests/GenerativeModelTests.swift index 2af18a7..3dc36e2 100644 --- a/Tests/GoogleAITests/GenerativeModelTests.swift +++ b/Tests/GoogleAITests/GenerativeModelTests.swift @@ -417,8 +417,8 @@ final class GenerativeModelTests: XCTestCase { for try await _ in stream { XCTFail("No content is there, this shouldn't happen.") } - } catch { - // TODO: Catch specific error. + } catch GenerateContentError.internalError(_ as InvalidCandidateError) { + // Underlying error is as expected, nothing else to check. return } @@ -440,8 +440,6 @@ final class GenerativeModelTests: XCTestCase { } catch let GenerateContentError.responseStoppedEarly(reason, _) { XCTAssertEqual(reason, .safety) return - } catch { - XCTFail("Wrong error generated: \(error)") } XCTFail("Should have caught an error.") @@ -548,7 +546,7 @@ final class GenerativeModelTests: XCTestCase { withExtension: "txt" ) - let stream = model.generateContentStream("Can you explain quantum physics?") + let stream = model.generateContentStream("Hi") var citations: [Citation] = [] for try await content in stream { XCTAssertNotNil(content.text) @@ -568,114 +566,67 @@ final class GenerativeModelTests: XCTestCase { func testGenerateContentStream_errorMidStream() async throws { MockURLProtocol.requestHandler = try httpRequestHandler( - forResource: "ExampleErrorMidStream", - withExtension: "json" + forResource: "streaming-failure-error-mid-stream", + withExtension: "txt" ) - let stream = model.generateContentStream("What sorts of questions can I ask you?") - - var textResponses = [String]() - var errorResponse: Error? + var responseCount = 0 do { + let stream = model.generateContentStream("Hi") for try await content in stream { XCTAssertNotNil(content.text) - let text = try XCTUnwrap(content.text) - textResponses.append(text) + responseCount += 1 } - } catch { - errorResponse = error + } catch let GenerateContentError.internalError(rpcError as RPCError) { + XCTAssertEqual(rpcError.httpResponseCode, 499) + XCTAssertEqual(rpcError.status, .cancelled) + + // Check the content count is correct. + XCTAssertEqual(responseCount, 2) + return } - // TODO: Add assertions for response content - XCTAssertEqual(textResponses.count, 2) - XCTAssertNotNil(errorResponse) + XCTFail("Expected an internalError with an RPCError.") } func testGenerateContentStream_nonHTTPResponse() async throws { MockURLProtocol.requestHandler = try nonHTTPRequestHandler() - let stream = model.generateContentStream("What sorts of questions can I ask you?") - var responseError: Error? + let stream = model.generateContentStream("Hi") do { for try await content in stream { XCTFail("Unexpected content in stream: \(content)") } - } catch { - responseError = error - } - - XCTAssertNotNil(responseError) - let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) - guard case let .internalError(underlyingError) = generateContentError else { - XCTFail("Not an internal error: \(generateContentError)") + } catch let GenerateContentError.internalError(underlying) { + XCTAssertEqual(underlying.localizedDescription, "Response was not an HTTP response.") return } - XCTAssertEqual(underlyingError.localizedDescription, "Response was not an HTTP response.") + + XCTFail("Expected an internal error.") } func testGenerateContentStream_invalidResponse() async throws { MockURLProtocol .requestHandler = try httpRequestHandler( - forResource: "InvalidStreamingResponse", - withExtension: "json" + forResource: "streaming-failure-invalid-json", + withExtension: "txt" ) let stream = model.generateContentStream(testPrompt) - var responseError: Error? do { for try await content in stream { XCTFail("Unexpected content in stream: \(content)") } - } catch { - responseError = error - } - - XCTAssertNotNil(responseError) - let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) - guard case let .internalError(underlyingError) = generateContentError else { - XCTFail("Not an internal error: \(generateContentError)") - return - } - let decodingError = try XCTUnwrap(underlyingError as? DecodingError) - guard case let .dataCorrupted(context) = decodingError else { - XCTFail("Not a data corrupted error: \(decodingError)") - return - } - XCTAssert(context.debugDescription.hasPrefix("Failed to decode GenerateContentResponse")) - } - - func testGenerateContentStream_emptyContent() async throws { - MockURLProtocol - .requestHandler = try httpRequestHandler( - forResource: "EmptyContentStreamingResponse", - withExtension: "json" - ) - - let stream = model.generateContentStream(testPrompt) - var responseError: Error? - do { - for try await content in stream { - XCTFail("Unexpected content in stream: \(content)") + } catch let GenerateContentError.internalError(underlying as DecodingError) { + guard case let .dataCorrupted(context) = underlying else { + XCTFail("Not a data corrupted error: \(underlying)") + return } - } catch { - responseError = error - } - - XCTAssertNotNil(responseError) - let generateContentError = try XCTUnwrap(responseError as? GenerateContentError) - guard case let .internalError(underlyingError) = generateContentError else { - XCTFail("Not an internal error: \(generateContentError)") - return - } - let invalidCandidateError = try XCTUnwrap(underlyingError as? InvalidCandidateError) - guard case let .emptyContent(emptyContentUnderlyingError) = invalidCandidateError else { - XCTFail("Not an empty content error: \(invalidCandidateError)") + XCTAssert(context.debugDescription.hasPrefix("Failed to decode GenerateContentResponse")) return } - _ = try XCTUnwrap( - emptyContentUnderlyingError as? DecodingError, - "Not a decoding error: \(emptyContentUnderlyingError)" - ) + + XCTFail("Expected an internal error.") } func testGenerateContentStream_malformedContent() async throws { diff --git a/Tests/GoogleAITests/SampleResponses/EmptyContentStreamingResponse.json b/Tests/GoogleAITests/SampleResponses/EmptyContentStreamingResponse.json deleted file mode 100644 index 175585e..0000000 --- a/Tests/GoogleAITests/SampleResponses/EmptyContentStreamingResponse.json +++ /dev/null @@ -1 +0,0 @@ -data: {"candidates": [{"content": {},"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "MEDIUM"},{"category": "HARM_CATEGORY_SEXUAL","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DEROGATORY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS","probability": "NEGLIGIBLE"}]}]} diff --git a/Tests/GoogleAITests/SampleResponses/ExampleErrorMidStream.json b/Tests/GoogleAITests/SampleResponses/ExampleErrorMidStream.json deleted file mode 100644 index c0471e8..0000000 --- a/Tests/GoogleAITests/SampleResponses/ExampleErrorMidStream.json +++ /dev/null @@ -1,17 +0,0 @@ -data: {"candidates": [{"content": {"parts": [{"text": " Sure, here are some questions you can ask me:\n\n* What is the weather like today?\n* What is the capital of France?\n*"}]},"index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}],"promptFeedback": {"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}} - -data: {"candidates": [{"content": {"parts": [{"text": " Who is the president of the United States?\n* What is the square root of 144?\n* What is the definition of the word \""}]},"index": 0,"safetyRatings": [{"category": "HARM_CATEGORY_TOXICITY","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_SEXUALLY_EXPLICIT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_VIOLENCE","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HATE_SPEECH","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_HARASSMENT","probability": "NEGLIGIBLE"},{"category": "HARM_CATEGORY_DANGEROUS_CONTENT","probability": "NEGLIGIBLE"}]}]} - -{ - "error": { - "code": 499, - "message": "The operation was cancelled.", - "status": "CANCELLED", - "details": [ - { - "@type": "type.googleapis.com/google.rpc.DebugInfo", - "detail": "[ORIGINAL ERROR] generic::cancelled: " - } - ] - } -}