From ba223a968fc49c66ef7211c266c7de3e4388fe6a Mon Sep 17 00:00:00 2001 From: Andrew Heard Date: Fri, 4 Oct 2024 18:45:29 -0400 Subject: [PATCH] Fix integration tests --- .../Tests/Integration/IntegrationTests.swift | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift index 8eddf79a648..5f5f08a5a6c 100644 --- a/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift +++ b/FirebaseVertexAI/Tests/Integration/IntegrationTests.swift @@ -96,12 +96,12 @@ final class IntegrationTests: XCTestCase { } func testCountTokens_image_fileData() async throws { - let fileData = ModelContent(parts: [.fileData( - mimetype: "image/jpeg", - uri: "gs://ios-opensource-samples.appspot.com/ios/public/blank.jpg" - )]) + let fileData = FileDataPart( + uri: "gs://ios-opensource-samples.appspot.com/ios/public/blank.jpg", + mimeType: "image/jpeg" + ) - let response = try await model.countTokens([fileData]) + let response = try await model.countTokens(fileData) XCTAssertEqual(response.totalTokens, 266) XCTAssertEqual(response.totalBillableCharacters, 35) @@ -122,9 +122,9 @@ final class IntegrationTests: XCTestCase { let sumResponse = FunctionResponse(name: "sum", response: ["result": .number(42)]) let response = try await model.countTokens([ - ModelContent(role: "user", parts: [.text(prompt)]), - ModelContent(role: "model", parts: [.functionCall(sumCall)]), - ModelContent(role: "function", parts: [.functionResponse(sumResponse)]), + ModelContent(role: "user", parts: prompt), + ModelContent(role: "model", parts: FunctionCallPart(sumCall)), + ModelContent(role: "function", parts: FunctionResponsePart(functionResponse: sumResponse)), ]) XCTAssertEqual(response.totalTokens, 24)