Skip to content

[AI] Unit tests for urlContext APIs #15343

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 11 commits into from
Sep 22, 2025
Merged

[AI] Unit tests for urlContext APIs #15343

Show file tree
Hide file tree
Changes from 3 commits
Commits
Show all changes
11 commits
Select commit Hold shift + click to select a range
538b890
checkpoint
paulb777 Sep 19, 2025
774ce1c
Tests for seven new mock responses now pass
paulb777 Sep 19, 2025
48eda93
style
paulb777 Sep 19, 2025
71210ea
Apply suggestions from code review
paulb777 Sep 19, 2025
1f410ae
improve function name
paulb777 Sep 19, 2025
ebd9444
empty urlMetadata array tests
paulb777 Sep 19, 2025
7df2377
Apply suggestions from code review
paulb777 Sep 19, 2025
83ee78a
[Firebase AI] Add `Candidate` decoding tests for `urlMetadata` (#15348)
andrewheard Sep 22, 2025
ada7b65
review
paulb777 Sep 22, 2025
729edba
cleanup
paulb777 Sep 22, 2025
8382596
Update FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
paulb777 Sep 22, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
68 changes: 68 additions & 0 deletions FirebaseAI/Tests/Unit/GenerativeModelGoogleAITests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,6 +337,53 @@ final class GenerativeModelGoogleAITests: XCTestCase {
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 160)
}

func testGenerateContent_success_urlContext() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-url-context",
withExtension: "json",
subdirectory: googleAISubdirectory
)

let response = try await model.generateContent(testPrompt)

XCTAssertEqual(response.candidates.count, 1)
let candidate = try XCTUnwrap(response.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
XCTAssertEqual(
retrievedURL,
URL(string: "https://berkshirehathaway.com")
)
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
let usageMetadata = try XCTUnwrap(response.usageMetadata)
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 424)
}

func testGenerateContent_success_urlContext_mixedValidity() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-url-context-mixed-validity",
withExtension: "json",
subdirectory: googleAISubdirectory
)

let response = try await model.generateContent(testPrompt)

let candidate = try XCTUnwrap(response.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)

let paywallURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[0])
XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)

let successURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[1])
XCTAssertEqual(successURLMetadata.retrievalStatus, .success)

let errorURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[2])
XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
}

func testGenerateContent_failure_invalidAPIKey() async throws {
let expectedStatusCode = 400
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
Expand Down Expand Up @@ -644,4 +691,25 @@ final class GenerativeModelGoogleAITests: XCTestCase {
let lastResponse = try XCTUnwrap(responses.last)
XCTAssertEqual(lastResponse.text, "text8")
}

func testGenerateContentStream_success_urlContext() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "streaming-success-url-context",
withExtension: "txt",
subdirectory: googleAISubdirectory
)

var responses = [GenerateContentResponse]()
let stream = try model.generateContentStream(testPrompt)
for try await response in stream {
responses.append(response)
}

let firstResponse = try XCTUnwrap(responses.first)
let candidate = try XCTUnwrap(firstResponse.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
}
}
86 changes: 86 additions & 0 deletions FirebaseAI/Tests/Unit/GenerativeModelVertexAITests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -491,6 +491,70 @@ final class GenerativeModelVertexAITests: XCTestCase {
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 371)
}

func testGenerateContent_success_urlContext() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-url-context",
withExtension: "json",
subdirectory: vertexSubdirectory
)

let response = try await model.generateContent(testPrompt)

XCTAssertEqual(response.candidates.count, 1)
let candidate = try XCTUnwrap(response.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
let retrievedURL = try XCTUnwrap(urlMetadata.retrievedURL)
XCTAssertEqual(
retrievedURL,
URL(string: "https://berkshirehathaway.com")
)
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
let usageMetadata = try XCTUnwrap(response.usageMetadata)
XCTAssertEqual(usageMetadata.toolUsePromptTokenCount, 34)
XCTAssertEqual(usageMetadata.thoughtsTokenCount, 36)
}

func testGenerateContent_success_urlContext_mixedValidity() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-url-context-mixed-validity",
withExtension: "json",
subdirectory: vertexSubdirectory
)

let response = try await model.generateContent(testPrompt)

let candidate = try XCTUnwrap(response.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 3)

let paywallURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[0])
XCTAssertEqual(paywallURLMetadata.retrievalStatus, .error)

let successURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[1])
XCTAssertEqual(successURLMetadata.retrievalStatus, .success)

let errorURLMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata[2])
XCTAssertEqual(errorURLMetadata.retrievalStatus, .error)
}

func testGenerateContent_success_urlContext_missingRetrievedURL() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-url-context-missing-retrievedurl",
withExtension: "json",
subdirectory: vertexSubdirectory
)

let response = try await model.generateContent(testPrompt)

let candidate = try XCTUnwrap(response.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
XCTAssertEqual(urlMetadata.retrievedURL?.absoluteString, "https://example.com/8")
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

There seems to be a contradiction in this test. The test name testGenerateContent_success_urlContext_missingRetrievedURL and the mock resource name unary-success-url-context-missing-retrievedurl suggest that the retrievedURL should be missing or nil. However, the assertion checks if retrievedURL?.absoluteString is equal to a specific string, which implies retrievedURL is not nil.

If the intention is to test for a missing URL, the assertion should probably be XCTAssertNil(urlMetadata.retrievedURL). If the current assertion is correct, consider renaming the test to avoid confusion.

Suggested change
XCTAssertEqual(urlMetadata.retrievedURL?.absoluteString, "https://example.com/8")
XCTAssertNil(urlMetadata.retrievedURL)

Copy link
Member Author

@paulb777 paulb777 Sep 19, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Renamed function to testGenerateContent_success_urlContext_retrievedURLPresentOnErrorStatus cc: @dlarocque @andrewheard

XCTAssertEqual(urlMetadata.retrievalStatus, .error)
}

func testGenerateContent_success_image_invalidSafetyRatingsIgnored() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "unary-success-image-invalid-safety-ratings",
Expand Down Expand Up @@ -1720,6 +1784,28 @@ final class GenerativeModelVertexAITests: XCTestCase {
XCTAssertEqual(responses, 1)
}

func testGenerateContentStream_success_urlContext() async throws {
MockURLProtocol.requestHandler = try GenerativeModelTestUtil.httpRequestHandler(
forResource: "streaming-success-url-context",
withExtension: "txt",
subdirectory: vertexSubdirectory
)

var responses = [GenerateContentResponse]()
let stream = try model.generateContentStream(testPrompt)
for try await response in stream {
responses.append(response)
}

let firstResponse = try XCTUnwrap(responses.first)
let candidate = try XCTUnwrap(firstResponse.candidates.first)
let urlContextMetadata = try XCTUnwrap(candidate.urlContextMetadata)
XCTAssertEqual(urlContextMetadata.urlMetadata.count, 1)
let urlMetadata = try XCTUnwrap(urlContextMetadata.urlMetadata.first)
XCTAssertEqual(urlMetadata.retrievedURL?.absoluteString, "https://google.com")
XCTAssertEqual(urlMetadata.retrievalStatus, .success)
}

// MARK: - Count Tokens

func testCountTokens_succeeds() async throws {
Expand Down
Loading