Skip to content

Commit 730e2b8

Browse files
committed
Add streaming generate content test
1 parent 14218a7 commit 730e2b8

File tree

2 files changed

+69
-0
lines changed

2 files changed

+69
-0
lines changed
Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
data: {"candidates": [{"content": {"parts": [{"text": "Thoughts"}],"role": "model"},"finishReason": "STOP","index": 0}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 1,"totalTokenCount": 22}}
2+
3+
data: {"candidates": [{"content": {"parts": [{"text": ": I can use the `print()` function in Python to print strings. "}],"role": "model"},"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"}]}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 16,"totalTokenCount": 37}}
4+
5+
data: {"candidates": [{"content": {"parts": [{"text": "\n\n"}],"role": "model"},"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"}]}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 16,"totalTokenCount": 37}}
6+
7+
data: {"candidates": [{"content": {"parts": [{"executableCode": {"language": "PYTHON","code": "\nprint(\"Hello, world!\")\n"}}],"role": "model"},"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"}]}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 29,"totalTokenCount": 50}}
8+
9+
data: {"candidates": [{"content": {"parts": [{"codeExecutionResult": {"outcome": "OUTCOME_OK","output": "Hello, world!\n"}}],"role": "model"},"index": 0}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 29,"totalTokenCount": 50}}
10+
11+
data: {"candidates": [{"content": {"parts": [{"text": "OK"}],"role": "model"},"finishReason": "STOP","index": 0}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 1,"totalTokenCount": 22}}
12+
13+
data: {"candidates": [{"content": {"parts": [{"text": ". I have printed \"Hello, world!\" using the `print()` function in"}],"role": "model"},"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"}]}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 17,"totalTokenCount": 38}}
14+
15+
data: {"candidates": [{"content": {"parts": [{"text": " Python. \n"}],"role": "model"},"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"}]}],"usageMetadata": {"promptTokenCount": 21,"candidatesTokenCount": 19,"totalTokenCount": 40}}
16+

Tests/GoogleAITests/GenerativeModelTests.swift

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -872,6 +872,59 @@ final class GenerativeModelTests: XCTestCase {
872872
}))
873873
}
874874

875+
func testGenerateContentStream_success_codeExecution() async throws {
876+
MockURLProtocol
877+
.requestHandler = try httpRequestHandler(
878+
forResource: "streaming-success-code-execution",
879+
withExtension: "txt"
880+
)
881+
let expectedTexts1 = [
882+
"Thoughts",
883+
": I can use the `print()` function in Python to print strings. ",
884+
"\n\n",
885+
]
886+
let expectedTexts2 = [
887+
"OK",
888+
". I have printed \"Hello, world!\" using the `print()` function in",
889+
" Python. \n",
890+
]
891+
let expectedTexts = Set(expectedTexts1 + expectedTexts2)
892+
let expectedLanguage = "PYTHON"
893+
let expectedCode = "\nprint(\"Hello, world!\")\n"
894+
let expectedOutput = "Hello, world!\n"
895+
896+
var textValues = [String]()
897+
let stream = model.generateContentStream(testPrompt)
898+
for try await content in stream {
899+
let candidate = try XCTUnwrap(content.candidates.first)
900+
let part = try XCTUnwrap(candidate.content.parts.first)
901+
switch part {
902+
case let .text(textPart):
903+
XCTAssertTrue(expectedTexts.contains(textPart))
904+
case let .executableCode(executableCode):
905+
XCTAssertEqual(executableCode.language, expectedLanguage)
906+
XCTAssertEqual(executableCode.code, expectedCode)
907+
case let .codeExecutionResult(codeExecutionResult):
908+
XCTAssertEqual(codeExecutionResult.outcome, .ok)
909+
XCTAssertEqual(codeExecutionResult.output, expectedOutput)
910+
default:
911+
XCTFail("Unexpected part type: \(part)")
912+
}
913+
try textValues.append(XCTUnwrap(content.text))
914+
}
915+
916+
XCTAssertEqual(textValues.joined(separator: "\n"), """
917+
\(expectedTexts1.joined(separator: "\n"))
918+
```\(expectedLanguage.lowercased())
919+
\(expectedCode)
920+
```
921+
```
922+
\(expectedOutput)
923+
```
924+
\(expectedTexts2.joined(separator: "\n"))
925+
""")
926+
}
927+
875928
func testGenerateContentStream_usageMetadata() async throws {
876929
MockURLProtocol
877930
.requestHandler = try httpRequestHandler(

0 commit comments

Comments
 (0)