Fix integration tests

firebase · Oct 4, 2024 · ba223a9 · ba223a9
1 parent 373bff4
commit ba223a9
Showing 1 changed file with 8 additions and 8 deletions.
diff --git 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)