From 231564f2368a9ccb2757525654209bf61c6611b7 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 26 Feb 2025 10:39:45 -0500 Subject: [PATCH 1/2] [Vertex AI] Add integration test for `generateContentStream` --- .../Tests/Integration/IntegrationTests.swift | 39 +++++++++++++++++++ 1 file changed, 39 insertions(+) diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift index 75a4a52e9f1..3989e868588 100644 --- a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift @@ -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 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() } From 442d57bc91d1014a2d4910a693c0b885eb8e2da2 Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Wed, 26 Feb 2025 10:57:34 -0500 Subject: [PATCH 2/2] Fix typo "from from" --- .../Tests/TestApp/Tests/Integration/IntegrationTests.swift | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift index 3989e868588..b5bfc94b93b 100644 --- a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift @@ -102,8 +102,8 @@ final class IntegrationTests: XCTestCase { 8. Neptune """ let prompt = """ - What are the names of the planets in the solar system, ordered from from closest to furthest - from the sun? Answer with a Markdown numbered list of the names and no other text. + 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()