Add toolUsePromptTokenCount and toolUsePromptTokensDetails tests

andrewheard · andrewheard · commit 98bb902e9636 · 2025-09-05T21:48:34.000-04:00
diff --git a/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift
@@ -333,6 +333,8 @@ final class GenerativeModelGoogleAITests: XCTestCase {
     let textPart = try XCTUnwrap(parts[2] as? TextPart)
     XCTAssertFalse(textPart.isThought)
     XCTAssertTrue(textPart.text.hasPrefix("The first 5 prime numbers are 2, 3, 5, 7, and 11."))
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)
   }
 
   func testGenerateContent_failure_invalidAPIKey() async throws {
diff --git a/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift b/FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
@@ -487,6 +487,8 @@ final class GenerativeModelVertexAITests: XCTestCase {
     XCTAssertEqual(
       textPart2.text, "The sum of the first 5 prime numbers (2, 3, 5, 7, and 11) is 28."
     )
+    let usageMetadata = try XCTUnwrap(response.usageMetadata)
+    XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371)
   }
 
   func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {

Original file line number	Diff line number	Diff line change
`@@ -333,6 +333,8 @@ final class GenerativeModelGoogleAITests: XCTestCase {`
`333`	`333`	`let textPart = try XCTUnwrap(parts[2] as? TextPart)`
`334`	`334`	`XCTAssertFalse(textPart.isThought)`
`335`	`335`	`XCTAssertTrue(textPart.text.hasPrefix("The first 5 prime numbers are 2, 3, 5, 7, and 11."))`
	`336`	`+ let usageMetadata = try XCTUnwrap(response.usageMetadata)`
	`337`	`+ XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)`
`336`	`338`	`}`
`337`	`339`
`338`	`340`	`func testGenerateContent_failure_invalidAPIKey() async throws {`
Original file line number	Diff line number	Diff line change
`@@ -487,6 +487,8 @@ final class GenerativeModelVertexAITests: XCTestCase {`
`487`	`487`	`XCTAssertEqual(`
`488`	`488`	`textPart2.text, "The sum of the first 5 prime numbers (2, 3, 5, 7, and 11) is 28."`
`489`	`489`	`)`
	`490`	`+ let usageMetadata = try XCTUnwrap(response.usageMetadata)`
	`491`	`+ XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371)`
`490`	`492`	`}`
`491`	`493`
`492`	`494`	`func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {`