Skip to content

Commit e02fd29

Browse files
paulb777gemini-code-assist[bot]andrewheard
committed
[AI] Unit tests for urlContext APIs (#15343)
Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: Andrew Heard <[email protected]>
1 parent 1fe4d88 commit e02fd29

File tree

4 files changed

+214
-3
lines changed

4 files changed

+214
-3
lines changed

‎FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift‎

Lines changed: 70 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -337,6 +337,53 @@ final class GenerativeModelGoogleAITests: XCTestCase {
337337
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)
338338
}
339339

340+
func testGenerateContent_success_urlContext() async throws {
341+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
342+
forResource: "unary-success-url-context",
343+
withExtension: "json",
344+
subdirectory: googleAISubdirectory
345+
)
346+
347+
let response = try await model.generateContent(testPrompt)
348+
349+
XCTAssertEqual(response.candidates.count, 1)
350+
let candidate = try XCTUnwrap(response.candidates.first)
351+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
352+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
353+
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
354+
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
355+
XCTAssertEqual(
356+
retrievedURL,
357+
URL(string: "https://berkshirehathaway.com")
358+
)
359+
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
360+
let usageMetadata = try XCTUnwrap(response.usageMetadata)
361+
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 424)
362+
}
363+
364+
func testGenerateContent_success_urlContext_mixedValidity() async throws {
365+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
366+
forResource: "unary-success-url-context-mixed-validity",
367+
withExtension: "json",
368+
subdirectory: googleAISubdirectory
369+
)
370+
371+
let response = try await model.generateContent(testPrompt)
372+
373+
let candidate = try XCTUnwrap(response.candidates.first)
374+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
375+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)
376+
377+
let paywallURLMetadata = urlContextMetadata.urlMetadata[0]
378+
XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)
379+
380+
let successURLMetadata = urlContextMetadata.urlMetadata[1]
381+
XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
382+
383+
let errorURLMetadata = urlContextMetadata.urlMetadata[2]
384+
XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
385+
}
386+
340387
func testGenerateContent_failure_invalidAPIKey() async throws {
341388
let expectedStatusCode = 400
342389
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
@@ -644,4 +691,27 @@ final class GenerativeModelGoogleAITests: XCTestCase {
644691
let lastResponse = try XCTUnwrap(responses.last)
645692
XCTAssertEqual(lastResponse.text, "text8")
646693
}
694+
695+
func testGenerateContentStream_success_urlContext() async throws {
696+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
697+
forResource: "streaming-success-url-context",
698+
withExtension: "txt",
699+
subdirectory: googleAISubdirectory
700+
)
701+
702+
var responses = [GenerateContentResponse]()
703+
let stream = try model.generateContentStream(testPrompt)
704+
for try await response in stream {
705+
responses.append(response)
706+
}
707+
708+
let firstResponse = try XCTUnwrap(responses.first)
709+
let candidate = try XCTUnwrap(firstResponse.candidates.first)
710+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
711+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
712+
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
713+
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
714+
XCTAssertEqual(retrievedURL, URL(string: "https://google.com"))
715+
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
716+
}
647717
}

‎FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift‎

Lines changed: 88 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -491,6 +491,71 @@ final class GenerativeModelVertexAITests: XCTestCase {
491491
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371)
492492
}
493493

494+
func testGenerateContent_success_urlContext() async throws {
495+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
496+
forResource: "unary-success-url-context",
497+
withExtension: "json",
498+
subdirectory: vertexSubdirectory
499+
)
500+
501+
let response = try await model.generateContent(testPrompt)
502+
503+
XCTAssertEqual(response.candidates.count, 1)
504+
let candidate = try XCTUnwrap(response.candidates.first)
505+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
506+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
507+
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
508+
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
509+
XCTAssertEqual(
510+
retrievedURL,
511+
URL(string: "https://berkshirehathaway.com")
512+
)
513+
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
514+
let usageMetadata = try XCTUnwrap(response.usageMetadata)
515+
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 34)
516+
XCTAssertEqual(usageMetadata.thoughtsTokenCount, 36)
517+
}
518+
519+
func testGenerateContent_success_urlContext_mixedValidity() async throws {
520+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
521+
forResource: "unary-success-url-context-mixed-validity",
522+
withExtension: "json",
523+
subdirectory: vertexSubdirectory
524+
)
525+
526+
let response = try await model.generateContent(testPrompt)
527+
528+
let candidate = try XCTUnwrap(response.candidates.first)
529+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
530+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)
531+
532+
let paywallURLMetadata = urlContextMetadata.urlMetadata[0]
533+
XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)
534+
535+
let successURLMetadata = urlContextMetadata.urlMetadata[1]
536+
XCTAssertEqual(successURLMetadata.retrievalStatus, .success)
537+
538+
let errorURLMetadata = urlContextMetadata.urlMetadata[2]
539+
XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
540+
}
541+
542+
func testGenerateContent_success_urlContext_retrievedURLPresentOnErrorStatus() async throws {
543+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
544+
forResource: "unary-success-url-context-missing-retrievedurl",
545+
withExtension: "json",
546+
subdirectory: vertexSubdirectory
547+
)
548+
549+
let response = try await model.generateContent(testPrompt)
550+
551+
let candidate = try XCTUnwrap(response.candidates.first)
552+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
553+
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
554+
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
555+
XCTAssertEqual(retrievedURL.absoluteString, "https://example.com/8")
556+
XCTAssertEqual(urlMetadata.retrievalStatus, .error)
557+
}
558+
494559
func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {
495560
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
496561
forResource: "unary-success-image-invalid-safety-ratings",
@@ -1720,6 +1785,29 @@ final class GenerativeModelVertexAITests: XCTestCase {
17201785
XCTAssertEqual(responses, 1)
17211786
}
17221787

1788+
func testGenerateContentStream_success_urlContext() async throws {
1789+
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
1790+
forResource: "streaming-success-url-context",
1791+
withExtension: "txt",
1792+
subdirectory: vertexSubdirectory
1793+
)
1794+
1795+
var responses = [GenerateContentResponse]()
1796+
let stream = try model.generateContentStream(testPrompt)
1797+
for try await response in stream {
1798+
responses.append(response)
1799+
}
1800+
1801+
let firstResponse = try XCTUnwrap(responses.first)
1802+
let candidate = try XCTUnwrap(firstResponse.candidates.first)
1803+
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
1804+
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
1805+
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
1806+
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
1807+
XCTAssertEqual(retrievedURL, URL(string: "https://google.com"))
1808+
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
1809+
}
1810+
17231811
// MARK: - Count Tokens
17241812

17251813
func testCountTokens_succeeds() async throws {

‎FirebaseAI/Tests/Unit/MockURLProtocol.swift‎

Lines changed: 5 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
2121
URLResponse,
2222
AsyncLineSequence<URL.AsyncBytes>?
2323
))?
24+
2425
override class func canInit(with request: URLRequest) -> Bool {
2526
#if os(watchOS)
2627
print("MockURLProtocol cannot be used on watchOS.")
@@ -33,13 +34,14 @@ class MockURLProtocol: URLProtocol, @unchecked Sendable {
3334
override class func canonicalRequest(for request: URLRequest) -> URLRequest { return request }
3435

3536
override func startLoading() {
36-
guard let requestHandler = MockURLProtocol.requestHandler else {
37-
fatalError("`requestHandler` is nil.")
38-
}
3937
guard let client = client else {
4038
fatalError("`client` is nil.")
4139
}
4240

41+
guard let requestHandler = MockURLProtocol.requestHandler else {
42+
fatalError("No request handler set.")
43+
}
44+
4345
Task {
4446
let (response, stream) = try requestHandler(self.request)
4547
client.urlProtocol(self, didReceive: response, cacheStoragePolicy: .notAllowed)

‎FirebaseAI/Tests/Unit/Types/GenerateContentResponseTests.swift‎

Lines changed: 51 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@ import XCTest
1717

1818
@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
1919
final class GenerateContentResponseTests: XCTestCase {
20+
let jsonDecoder = JSONDecoder()
21+
2022
// MARK: - GenerateContentResponse Computed Properties
2123

2224
func testGenerateContentResponse_inlineDataParts_success() throws {
@@ -106,4 +108,53 @@ final class GenerateContentResponseTests: XCTestCase {
106108
"functionCalls should be empty when there are no candidates."
107109
)
108110
}
111+
112+
// MARK: - Decoding Tests
113+
114+
func testDecodeCandidate_emptyURLMetadata_urlContextMetadataIsNil() throws {
115+
let json = """
116+
{
117+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
118+
"finishReason": "STOP",
119+
"urlContextMetadata": { "urlMetadata": [] }
120+
}
121+
"""
122+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
123+
124+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
125+
126+
XCTAssertNil(
127+
candidate.urlContextMetadata,
128+
"urlContextMetadata should be nil if the `urlMetadata` array is empty in the candidate."
129+
)
130+
XCTAssertEqual(candidate.content.role, "model")
131+
let part = try XCTUnwrap(candidate.content.parts.first)
132+
let textPart = try XCTUnwrap(part as? TextPart)
133+
XCTAssertEqual(textPart.text, "Some text.")
134+
XCTAssertFalse(textPart.isThought)
135+
XCTAssertEqual(candidate.finishReason, .stop)
136+
}
137+
138+
func testDecodeCandidate_missingURLMetadata_urlContextMetadataIsNil() throws {
139+
let json = """
140+
{
141+
"content": { "role": "model", "parts": [ { "text": "Some text." } ] },
142+
"finishReason": "STOP"
143+
}
144+
"""
145+
let jsonData = try XCTUnwrap(json.data(using: .utf8))
146+
147+
let candidate = try jsonDecoder.decode(Candidate.self, from: jsonData)
148+
149+
XCTAssertNil(
150+
candidate.urlContextMetadata,
151+
"urlContextMetadata should be nil if `urlMetadata` is not provided in the candidate."
152+
)
153+
XCTAssertEqual(candidate.content.role, "model")
154+
let part = try XCTUnwrap(candidate.content.parts.first)
155+
let textPart = try XCTUnwrap(part as? TextPart)
156+
XCTAssertEqual(textPart.text, "Some text.")
157+
XCTAssertFalse(textPart.isThought)
158+
XCTAssertEqual(candidate.finishReason, .stop)
159+
}
109160
}

0 commit comments

Comments
 (0)