Skip to content

[Vertex AI] Add integration test for generateContentStream #14501

New issue

Have a question about this project? # for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “#”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? # to your account

Merged
merged 2 commits into from
Feb 26, 2025
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -90,6 +90,45 @@ final class IntegrationTests: XCTestCase {
XCTAssertEqual(candidatesTokensDetails.tokenCount, usageMetadata.candidatesTokenCount)
}

func testGenerateContentStream() async throws {
let expectedText = """
1. Mercury
2. Venus
3. Earth
4. Mars
5. Jupiter
6. Saturn
7. Uranus
8. Neptune
"""
let prompt = """
What are the names of the planets in the solar system, ordered from closest to furthest from
the sun? Answer with a Markdown numbered list of the names and no other text.
"""
let chat = model.startChat()

let stream = try chat.sendMessageStream(prompt)
var textValues = [String]()
for try await value in stream {
try textValues.append(XCTUnwrap(value.text))
}

let userHistory = try XCTUnwrap(chat.history.first)
XCTAssertEqual(userHistory.role, "user")
XCTAssertEqual(userHistory.parts.count, 1)
let promptTextPart = try XCTUnwrap(userHistory.parts.first as? TextPart)
XCTAssertEqual(promptTextPart.text, prompt)
let modelHistory = try XCTUnwrap(chat.history.last)
XCTAssertEqual(modelHistory.role, "model")
XCTAssertEqual(modelHistory.parts.count, 1)
let modelTextPart = try XCTUnwrap(modelHistory.parts.first as? TextPart)
let modelText = modelTextPart.text.trimmingCharacters(in: .whitespacesAndNewlines)
XCTAssertEqual(modelText, expectedText)
XCTAssertGreaterThan(textValues.count, 1)
let text = textValues.joined().trimmingCharacters(in: .whitespacesAndNewlines)
XCTAssertEqual(text, expectedText)
}

func testGenerateContent_appCheckNotConfigured_shouldFail() async throws {
let app = try FirebaseApp.defaultNamedCopy(name: TestAppCheckProviderFactory.notConfiguredName)
addTeardownBlock { await app.delete() }
Expand Down
Loading