Skip to content

Commit 09dcec1

Browse files
authored
[AI] Add an implicit caching unit test (#15727)
1 parent 8b29daf commit 09dcec1

File tree

1 file changed

+82
-0
lines changed

1 file changed

+82
-0
lines changed
Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
// Copyright 2026 Google LLC
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
import FirebaseAppCheckInterop
16+
import FirebaseAuthInterop
17+
import FirebaseCore
18+
import XCTest
19+
20+
@testable import FirebaseAILogic
21+
22+
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
23+
final class GenerativeModelImplicitCachingTests: XCTestCase {
24+
let testPrompt = "What sorts of questions can I ask you?"
25+
let testModelName = "test-model"
26+
let testModelResourceName =
27+
"projects/test-project-id/locations/test-location/publishers/google/models/test-model"
28+
let apiConfig = FirebaseAI.defaultVertexAIAPIConfig
29+
30+
let vertexSubdirectory = "mock-responses/vertexai"
31+
32+
var urlSession: URLSession!
33+
var model: GenerativeModel!
34+
35+
override func setUp() async throws {
36+
let configuration = URLSessionConfiguration.default
37+
configuration.protocolClasses = [MockURLProtocol.self]
38+
urlSession = try XCTUnwrap(URLSession(configuration: configuration))
39+
model = GenerativeModel(
40+
modelName: testModelName,
41+
modelResourceName: testModelResourceName,
42+
firebaseInfo: GenerativeModelTestUtil.testFirebaseInfo(),
43+
apiConfig: apiConfig,
44+
tools: nil,
45+
requestOptions: RequestOptions(),
46+
urlSession: urlSession
47+
)
48+
}
49+
50+
override func tearDown() {
51+
MockURLProtocol.requestHandler = nil
52+
}
53+
54+
func testGenerateContent_success_implicitCaching() async throws {
55+
MockURLProtocol
56+
.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
57+
forResource: "unary-success-implicit-caching",
58+
withExtension: "json",
59+
subdirectory: vertexSubdirectory
60+
)
61+
62+
let response = try await model.generateContent(testPrompt)
63+
64+
XCTAssertEqual(response.candidates.count, 1)
65+
let candidate = try XCTUnwrap(response.candidates.first)
66+
let finishReason = try XCTUnwrap(candidate.finishReason)
67+
XCTAssertEqual(finishReason, .stop)
68+
69+
let usageMetadata = try XCTUnwrap(response.usageMetadata)
70+
XCTAssertEqual(usageMetadata.promptTokenCount, 12013)
71+
XCTAssertEqual(usageMetadata.candidatesTokenCount, 15)
72+
XCTAssertEqual(usageMetadata.totalTokenCount, 12101)
73+
74+
// Validate implicit caching fields
75+
XCTAssertEqual(usageMetadata.cachedContentTokenCount, 11243)
76+
XCTAssertEqual(usageMetadata.promptTokensDetails.count, 1)
77+
let detail = try XCTUnwrap(usageMetadata.promptTokensDetails.first)
78+
XCTAssertEqual(detail.modality, .text)
79+
XCTAssertEqual(detail.tokenCount, 12013)
80+
XCTAssertEqual(usageMetadata.thoughtsTokenCount, 73)
81+
}
82+
}

0 commit comments

Comments
 (0)