Skip to content

Commit ba223a9

Browse files
committed
Fix integration tests
1 parent 373bff4 commit ba223a9

File tree

1 file changed

+8
-8
lines changed

1 file changed

+8
-8
lines changed

FirebaseVertexAI/Tests/Integration/IntegrationTests.swift

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -96,12 +96,12 @@ final class IntegrationTests: XCTestCase {
9696
}
9797

9898
func testCountTokens_image_fileData() async throws {
99-
let fileData = ModelContent(parts: [.fileData(
100-
mimetype: "image/jpeg",
101-
uri: "gs://ios-opensource-samples.appspot.com/ios/public/blank.jpg"
102-
)])
99+
let fileData = FileDataPart(
100+
uri: "gs://ios-opensource-samples.appspot.com/ios/public/blank.jpg",
101+
mimeType: "image/jpeg"
102+
)
103103

104-
let response = try await model.countTokens([fileData])
104+
let response = try await model.countTokens(fileData)
105105

106106
XCTAssertEqual(response.totalTokens, 266)
107107
XCTAssertEqual(response.totalBillableCharacters, 35)
@@ -122,9 +122,9 @@ final class IntegrationTests: XCTestCase {
122122
let sumResponse = FunctionResponse(name: "sum", response: ["result": .number(42)])
123123

124124
let response = try await model.countTokens([
125-
ModelContent(role: "user", parts: [.text(prompt)]),
126-
ModelContent(role: "model", parts: [.functionCall(sumCall)]),
127-
ModelContent(role: "function", parts: [.functionResponse(sumResponse)]),
125+
ModelContent(role: "user", parts: prompt),
126+
ModelContent(role: "model", parts: FunctionCallPart(sumCall)),
127+
ModelContent(role: "function", parts: FunctionResponsePart(functionResponse: sumResponse)),
128128
])
129129

130130
XCTAssertEqual(response.totalTokens, 24)

0 commit comments

Comments
 (0)