@@ -39,6 +39,8 @@ final class IntegrationTests: XCTestCase {
39
39
SafetySetting ( harmCategory: . dangerousContent, threshold: . blockLowAndAbove) ,
40
40
SafetySetting ( harmCategory: . civicIntegrity, threshold: . blockLowAndAbove) ,
41
41
]
42
+ // Candidates and total token counts may differ slightly between runs due to whitespace tokens.
43
+ let tokenCountAccuracy = 1
42
44
43
45
var vertex : VertexAI !
44
46
var model : GenerativeModel !
@@ -74,6 +76,18 @@ final class IntegrationTests: XCTestCase {
74
76
75
77
let text = try XCTUnwrap ( response. text) . trimmingCharacters ( in: . whitespacesAndNewlines)
76
78
XCTAssertEqual ( text, " Mountain View " )
79
+ let usageMetadata = try XCTUnwrap ( response. usageMetadata)
80
+ XCTAssertEqual ( usageMetadata. promptTokenCount, 21 )
81
+ XCTAssertEqual ( usageMetadata. candidatesTokenCount, 3 , accuracy: tokenCountAccuracy)
82
+ XCTAssertEqual ( usageMetadata. totalTokenCount, 24 , accuracy: tokenCountAccuracy)
83
+ XCTAssertEqual ( usageMetadata. promptTokensDetails. count, 1 )
84
+ let promptTokensDetails = try XCTUnwrap ( usageMetadata. promptTokensDetails. first)
85
+ XCTAssertEqual ( promptTokensDetails. modality, . text)
86
+ XCTAssertEqual ( promptTokensDetails. tokenCount, usageMetadata. promptTokenCount)
87
+ XCTAssertEqual ( usageMetadata. candidatesTokensDetails. count, 1 )
88
+ let candidatesTokensDetails = try XCTUnwrap ( usageMetadata. candidatesTokensDetails. first)
89
+ XCTAssertEqual ( candidatesTokensDetails. modality, . text)
90
+ XCTAssertEqual ( candidatesTokensDetails. tokenCount, usageMetadata. candidatesTokenCount)
77
91
}
78
92
79
93
func testGenerateContent_appCheckNotConfigured_shouldFail( ) async throws {
@@ -113,6 +127,10 @@ final class IntegrationTests: XCTestCase {
113
127
114
128
XCTAssertEqual ( response. totalTokens, 14 )
115
129
XCTAssertEqual ( response. totalBillableCharacters, 51 )
130
+ XCTAssertEqual ( response. promptTokensDetails. count, 1 )
131
+ let promptTokensDetails = try XCTUnwrap ( response. promptTokensDetails. first)
132
+ XCTAssertEqual ( promptTokensDetails. modality, . text)
133
+ XCTAssertEqual ( promptTokensDetails. tokenCount, 14 )
116
134
}
117
135
118
136
#if canImport(UIKit)
@@ -126,6 +144,15 @@ final class IntegrationTests: XCTestCase {
126
144
127
145
XCTAssertEqual ( response. totalTokens, 266 )
128
146
XCTAssertEqual ( response. totalBillableCharacters, 35 )
147
+ XCTAssertEqual ( response. promptTokensDetails. count, 2 ) // Image prompt + system instruction
148
+ let textPromptTokensDetails = try XCTUnwrap ( response. promptTokensDetails. first {
149
+ $0. modality == . text
150
+ } ) // System instruction
151
+ XCTAssertEqual ( textPromptTokensDetails. tokenCount, 8 )
152
+ let imagePromptTokenDetails = try XCTUnwrap ( response. promptTokensDetails. first {
153
+ $0. modality == . image
154
+ } )
155
+ XCTAssertEqual ( imagePromptTokenDetails. tokenCount, 258 )
129
156
}
130
157
#endif // canImport(UIKit)
131
158
@@ -137,6 +164,15 @@ final class IntegrationTests: XCTestCase {
137
164
138
165
XCTAssertEqual ( response. totalTokens, 266 )
139
166
XCTAssertEqual ( response. totalBillableCharacters, 35 )
167
+ XCTAssertEqual ( response. promptTokensDetails. count, 2 ) // Image prompt + system instruction
168
+ let textPromptTokensDetails = try XCTUnwrap ( response. promptTokensDetails. first {
169
+ $0. modality == . text
170
+ } ) // System instruction
171
+ XCTAssertEqual ( textPromptTokensDetails. tokenCount, 8 )
172
+ let imagePromptTokenDetails = try XCTUnwrap ( response. promptTokensDetails. first {
173
+ $0. modality == . image
174
+ } )
175
+ XCTAssertEqual ( imagePromptTokenDetails. tokenCount, 258 )
140
176
}
141
177
142
178
func testCountTokens_image_fileData_requiresAuth_signedIn( ) async throws {
@@ -199,6 +235,10 @@ final class IntegrationTests: XCTestCase {
199
235
200
236
XCTAssertEqual ( response. totalTokens, 24 )
201
237
XCTAssertEqual ( response. totalBillableCharacters, 71 )
238
+ XCTAssertEqual ( response. promptTokensDetails. count, 1 )
239
+ let promptTokensDetails = try XCTUnwrap ( response. promptTokensDetails. first)
240
+ XCTAssertEqual ( promptTokensDetails. modality, . text)
241
+ XCTAssertEqual ( promptTokensDetails. tokenCount, 24 )
202
242
}
203
243
204
244
func testCountTokens_jsonSchema( ) async throws {
@@ -220,6 +260,10 @@ final class IntegrationTests: XCTestCase {
220
260
221
261
XCTAssertEqual ( response. totalTokens, 58 )
222
262
XCTAssertEqual ( response. totalBillableCharacters, 160 )
263
+ XCTAssertEqual ( response. promptTokensDetails. count, 1 )
264
+ let promptTokensDetails = try XCTUnwrap ( response. promptTokensDetails. first)
265
+ XCTAssertEqual ( promptTokensDetails. modality, . text)
266
+ XCTAssertEqual ( promptTokensDetails. tokenCount, 58 )
223
267
}
224
268
225
269
func testCountTokens_appCheckNotConfigured_shouldFail( ) async throws {
0 commit comments