diff --git a/Package.swift b/Package.swift index 1703f6b6486..827b6bc979a 100644 --- a/Package.swift +++ b/Package.swift @@ -39,7 +39,7 @@ let package = Package( .target( name: "SharedTestUtilities", path: "Tests/SharedTestUtilities", - swiftSettings: defaultSwiftSettings + swiftSettings: defaultSwiftSettings, ), .testTarget( name: "GeminiForFoundationModelsTests", @@ -49,6 +49,24 @@ let package = Package( ], swiftSettings: defaultSwiftSettings ), + .target( + name: "GeminiAPIClient", + dependencies: [ + "GenerateContentDataModels", + "HTTPStreamingClient", + "SharedDataModels", + ], + swiftSettings: defaultSwiftSettings + ), + .testTarget( + name: "GeminiAPIClientTests", + dependencies: [ + "GeminiAPIClient", + "HTTPStreamingClient", + "SharedTestUtilities", + ], + swiftSettings: defaultSwiftSettings + ), .target( name: "SharedDataModels", swiftSettings: defaultSwiftSettings diff --git a/Sources/GeminiAPIClient/GeminiAPIClient.swift b/Sources/GeminiAPIClient/GeminiAPIClient.swift new file mode 100644 index 00000000000..24b11b3723a --- /dev/null +++ b/Sources/GeminiAPIClient/GeminiAPIClient.swift @@ -0,0 +1,313 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package import Foundation +package import GenerateContentDataModels +import HTTPStreamingClient +import SharedDataModels + +#if canImport(FoundationNetworking) + package import FoundationNetworking +#endif + +// MARK: - Gemini API Client + +/// A client for communicating with Google Gemini backend endpoints. +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +package struct GeminiAPIClient: Sendable { + /// The resource path of the target model (e.g., `"v1beta/models/gemini-3.5-flash-lite"` + /// or `"v1beta1/projects/p/locations/l/publishers/google/models/gemini-3.5-flash-lite"`). + let modelResourcePath: String + + /// The base URL of the Gemini API endpoint. + let baseURL: URL + + /// An optional async provider for dynamic headers (such as API keys or Bearer tokens). + let headerProvider: (@Sendable () async throws -> [String: String])? + + private let httpClient: HTTPStreamingClient + + /// Initializes a new Gemini API client with a model resource path and target base URL. + /// + /// - Parameters: + /// - modelResourcePath: The resource path of the target model (e.g., + /// `"v1beta/models/gemini-3.5-flash-lite"`). + /// - baseURL: The base URL of the Gemini API endpoint. + /// - headerProvider: An optional async provider for dynamic headers (such as API keys or Bearer + /// tokens). + /// - sessionConfiguration: The `URLSessionConfiguration` to use. Defaults to `.ephemeral`. + package init( + modelResourcePath: String, + baseURL: URL, + headerProvider: (@Sendable () async throws -> [String: String])? = nil, + sessionConfiguration: URLSessionConfiguration = .ephemeral + ) { + assert(!modelResourcePath.isEmpty, "modelResourcePath must not be empty.") + self.modelResourcePath = modelResourcePath + self.baseURL = baseURL + self.headerProvider = headerProvider + self.httpClient = HTTPStreamingClient(configuration: sessionConfiguration) + } + + /// Sends a streaming text generation request and delivers responses asynchronously as a + /// backpressured `GenerateContentStream` sequence. + /// + /// - Parameter request: The structured content generation request. + /// - Returns: A backpressured `GenerateContentStream` async sequence. + /// - Throws: `GeminiAPIError.apiError` on API failures, or standard network errors. + package func generateContentStream( + for request: GenerateContentRequest + ) async throws -> GenerateContentStream { + let urlRequest = try await makeURLRequest( + action: "streamGenerateContent", + queryItems: [URLQueryItem(name: "alt", value: "sse")], + body: request + ) + + let (lines, response) = try await httpClient.lines(for: urlRequest) + + if response.statusCode != 200 { + let bodyData = try await collectBody(from: lines) + throw parseError(from: bodyData, statusCode: response.statusCode, response: response) + } + + return GenerateContentStream(lines: lines, response: response) + } + + /// Counts the number of tokens in the given request. + /// + /// - Parameter request: The token count calculation request. + /// - Returns: The calculated `CountTokensResponse`. + /// - Throws: `GeminiAPIError.apiError` on API failures, or standard network errors. + package func countTokens( + for request: CountTokensRequest + ) async throws -> CountTokensResponse { + let urlRequest = try await makeURLRequest( + action: "countTokens", + body: request + ) + + let (lines, response) = try await httpClient.lines(for: urlRequest) + let bodyData = try await collectBody(from: lines) + + if response.statusCode != 200 { + throw parseError(from: bodyData, statusCode: response.statusCode, response: response) + } + + return try JSONDecoder().decode(CountTokensResponse.self, from: bodyData) + } + + private func makeURLRequest( + action: String, + queryItems: [URLQueryItem]? = nil, + body: Body + ) async throws -> URLRequest { + guard var components = URLComponents(url: baseURL, resolvingAgainstBaseURL: false) else { + throw URLError(.badURL) + } + + let sanitizedResourcePath = modelResourcePath.drop { $0 == "/" } + let basePath = + components.path.hasSuffix("/") + ? String(components.path.dropLast()) + : components.path + components.path = "\(basePath)/\(sanitizedResourcePath):\(action)" + if let queryItems { + components.queryItems = (components.queryItems ?? []) + queryItems + } + guard let requestURL = components.url else { + throw URLError(.badURL) + } + + var urlRequest = URLRequest(url: requestURL) + urlRequest.httpMethod = "POST" + urlRequest.setValue("application/json", forHTTPHeaderField: "Content-Type") + + if let headerProvider { + for (key, value) in try await headerProvider() { + urlRequest.setValue(value, forHTTPHeaderField: key) + } + } + + urlRequest.httpBody = try JSONEncoder().encode(body) + return urlRequest + } + + private func collectBody(from lines: HTTPAsyncLineSequence) async throws -> Data { + return try await lines.reduce(into: Data()) { result, line in + if !result.isEmpty { + result.append(0x0A) // "\n" + } + result.append(contentsOf: line.utf8) + } + } +} + +// MARK: - Generate Content Stream + +/// An asynchronous sequence of `GenerateContentResponse` chunks streamed from Gemini. +/// +/// Iterates on-demand over Server-Sent Events with backpressure and zero unstructured `Task` +/// allocation. Cancellation propagates directly to the underlying network stream. +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +package struct GenerateContentStream: AsyncSequence, Sendable { + /// The type of element produced by this asynchronous sequence. + package typealias Element = GenerateContentResponse + + private let lines: HTTPAsyncLineSequence + private let response: HTTPURLResponse + + /// Initializes a new content stream from an underlying HTTP line sequence and response metadata. + /// + /// - Parameters: + /// - lines: The sequence of text lines received from the server. + /// - response: The initial HTTP response headers and status code. + init(lines: HTTPAsyncLineSequence, response: HTTPURLResponse) { + self.lines = lines + self.response = response + } + + /// Creates an asynchronous iterator over the stream of generated content response chunks. + /// + /// - Returns: An `AsyncIterator` instance. + package func makeAsyncIterator() -> AsyncIterator { + AsyncIterator(linesIterator: lines.makeAsyncIterator(), response: response) + } + + /// An asynchronous iterator over Server-Sent Events decoded into + /// `GenerateContentResponse` chunks. + package struct AsyncIterator: AsyncIteratorProtocol { + private var linesIterator: HTTPAsyncLineSequence.AsyncIterator + private let response: HTTPURLResponse + private let decoder = JSONDecoder() + private var sseDataBuffer = "" + private var extraLinesBuffer = "" + + init(linesIterator: HTTPAsyncLineSequence.AsyncIterator, response: HTTPURLResponse) { + self.linesIterator = linesIterator + self.response = response + } + + /// Asynchronously advances to and returns the next `GenerateContentResponse` chunk. + /// + /// - Returns: The next decoded `GenerateContentResponse`, or `nil` if the stream has finished. + /// - Throws: An error if reading or decoding fails, or if a mid-stream API error occurs. + package mutating func next() async throws -> GenerateContentResponse? { + while let line = try await linesIterator.next() { + // Empty line marks the end of an SSE event + if line.isEmpty || line.allSatisfy({ $0.isWhitespace }) { + if !sseDataBuffer.isEmpty { + let dataString = sseDataBuffer + sseDataBuffer = "" + return try decodeEventData(dataString) + } + continue + } + + // SSE comment line (e.g. ": keep-alive") + if line.hasPrefix(":") { + continue + } + + // SSE control fields (e.g. "event: message", "id: 1", "retry: 5000") + if line.hasPrefix("event:") || line.hasPrefix("id:") || line.hasPrefix("retry:") { + continue + } + + // SSE data field + if line.hasPrefix("data:") { + var dataContent = line.dropFirst(5) + // The SSE specification only allows a single leading space after the colon. + if dataContent.hasPrefix(" ") { + dataContent = dataContent.dropFirst() + } + if !dataContent.isEmpty { + if !sseDataBuffer.isEmpty { + sseDataBuffer.append("\n") + } + sseDataBuffer.append(contentsOf: dataContent) + } + continue + } + + // Non-SSE payload line (e.g. raw JSON error block or unexpected content) + extraLinesBuffer.append(line) + extraLinesBuffer.append("\n") + } + + // Flush any pending SSE event data + if !sseDataBuffer.isEmpty { + let dataString = sseDataBuffer + sseDataBuffer = "" + return try decodeEventData(dataString) + } + + // If extra non-SSE lines were accumulated, parse as error + let trimmedExtra = extraLinesBuffer.trimmingCharacters(in: .whitespacesAndNewlines) + if !trimmedExtra.isEmpty { + let data = Data(trimmedExtra.utf8) + throw parseError(from: data, statusCode: response.statusCode, response: response) + } + + return nil + } + + private func decodeEventData(_ jsonString: String) throws -> GenerateContentResponse { + let data = Data(jsonString.utf8) + // Fast-path: only attempt error decoding if the payload contains an "error" key. + // GoogleCloudAPIError requires top-level code and message, avoiding false positives. + if jsonString.contains("\"error\""), + let apiError = try? decoder.decode(GoogleCloudAPIError.self, from: data) + { + let headerRetryAfter = parseRetryAfterHeader(from: response) + let resolvedError = headerRetryAfter.map { apiError.withRetryDelay($0) } ?? apiError + throw GeminiAPIError.apiError(resolvedError) + } + return try decoder.decode(GenerateContentResponse.self, from: data) + } + } +} + +// MARK: - Internal Error Parsing Helpers + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private func parseError( + from data: Data, + statusCode: Int, + response: HTTPURLResponse +) -> GeminiAPIError { + if let apiError = try? JSONDecoder().decode(GoogleCloudAPIError.self, from: data) { + let headerRetryAfter = parseRetryAfterHeader(from: response) + let resolvedError = headerRetryAfter.map { apiError.withRetryDelay($0) } ?? apiError + return GeminiAPIError.apiError(resolvedError) + } else { + return GeminiAPIError.httpError( + statusCode: statusCode, + body: String(decoding: data, as: UTF8.self) + ) + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +private func parseRetryAfterHeader(from response: HTTPURLResponse) -> Duration? { + if let headerValue = response.value(forHTTPHeaderField: "Retry-After")?.trimmingCharacters( + in: .whitespaces + ) { + if let seconds = Double(headerValue), seconds >= 0 { + return .seconds(seconds) + } + } + + return nil +} diff --git a/Sources/GeminiAPIClient/GeminiAPIError.swift b/Sources/GeminiAPIClient/GeminiAPIError.swift new file mode 100644 index 00000000000..625c69fefe9 --- /dev/null +++ b/Sources/GeminiAPIClient/GeminiAPIError.swift @@ -0,0 +1,99 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +package import SharedDataModels + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +/// Errors thrown by `GeminiAPIClient`. +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +package enum GeminiAPIError: Error, Sendable, Equatable { + /// An API error returned by the Google Gemini service conforming to AIP-0193. + case apiError(GoogleCloudAPIError) + + /// An HTTP failure status code with a raw response body. + case httpError(statusCode: Int, body: String) +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +extension GeminiAPIError { + /// The retry delay advice for this error, if available from headers or payload details. + package var retryAfter: Duration? { + switch self { + case .apiError(let error): + error.retryDelay + case .httpError: + nil + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +extension GeminiAPIError: LocalizedError { + package static var errorDomain: String { "GeminiAPIClient.GeminiAPIError" } + + package var errorDescription: String? { + switch self { + case .apiError(let error): + error.errorDescription + case .httpError(let statusCode, let body): + "HTTP \(statusCode): \(body)" + } + } + + package var failureReason: String? { + switch self { + case .apiError(let error): + error.failureReason + case .httpError(let statusCode, _): + "HTTP status \(statusCode)" + } + } + + package var helpAnchor: String? { + switch self { + case .apiError(let error): + error.helpAnchor + case .httpError: + nil + } + } +} + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +extension GeminiAPIError: CustomNSError { + package var errorCode: Int { + switch self { + case .apiError(let error): + error.errorCode + case .httpError(let statusCode, _): + statusCode + } + } + + package var errorUserInfo: [String: Any] { + switch self { + case .apiError(let error): + error.errorUserInfo + case .httpError(let statusCode, let body): + [ + "statusCode": statusCode, + "body": body, + ] + } + } +} diff --git a/Sources/SharedDataModels/GoogleCloudAPIError.swift b/Sources/SharedDataModels/GoogleCloudAPIError.swift index 6a1d0aef964..d3d4ac2891e 100644 --- a/Sources/SharedDataModels/GoogleCloudAPIError.swift +++ b/Sources/SharedDataModels/GoogleCloudAPIError.swift @@ -14,20 +14,37 @@ import Foundation -/// Represents a Google Cloud API error response body as defined by AIP-0193. +/// Represents a Google Cloud API error response body as defined by AIP-193. @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) package struct GoogleCloudAPIError: Codable, Sendable, Equatable, Hashable { + /// Represents the nested error `Status` in a Google Cloud API error. + /// + /// See [AIP-193](https://google.aip.dev/193#http11json-representation) for more details. + private struct Status: Codable, Sendable, Equatable, Hashable { + /// The HTTP status code value. + let code: Int + + /// A developer-facing, human-readable English error message. + let message: String + + /// The canonical status code indicating the nature of the error. + let status: RPCErrorStatus? + + /// Additional details about the error. + let details: [Detail]? + } + /// The HTTP status code value. - package let code: Int + package var code: Int { error.code } /// A developer-facing, human-readable English error message. - package let message: String + package var message: String { error.message } /// The canonical status code indicating the nature of the error. - package let status: RPCErrorStatus? + package var status: RPCErrorStatus? { error.status } /// Additional details about the error. - package let details: [Detail]? + package var details: [Detail]? { error.details } /// The retry delay duration, if provided by `RetryInfo` details or HTTP headers. package var retryDelay: Duration? { @@ -44,19 +61,9 @@ package struct GoogleCloudAPIError: Codable, Sendable, Equatable, Hashable { return nil } + private let error: Status private let explicitRetryDelay: Duration? - private enum EnvelopeCodingKeys: String, CodingKey { - case error - } - - private enum StatusCodingKeys: String, CodingKey { - case code - case message - case status - case details - } - /// Creates a new `GoogleCloudAPIError`. /// /// - Parameters: @@ -72,44 +79,10 @@ package struct GoogleCloudAPIError: Codable, Sendable, Equatable, Hashable { details: [Detail]? = nil, retryDelay: Duration? = nil ) { - self.code = code - self.message = message - self.status = status - self.details = details + self.error = Status(code: code, message: message, status: status, details: details) self.explicitRetryDelay = retryDelay } - package init(from decoder: any Decoder) throws { - if let envelopeContainer = try? decoder.container(keyedBy: EnvelopeCodingKeys.self), - envelopeContainer.contains(.error) - { - let statusContainer = try envelopeContainer.nestedContainer( - keyedBy: StatusCodingKeys.self, forKey: .error) - self.code = try statusContainer.decode(Int.self, forKey: .code) - self.message = try statusContainer.decode(String.self, forKey: .message) - self.status = try statusContainer.decodeIfPresent(RPCErrorStatus.self, forKey: .status) - self.details = try statusContainer.decodeIfPresent([Detail].self, forKey: .details) - self.explicitRetryDelay = nil - } else { - let container = try decoder.container(keyedBy: StatusCodingKeys.self) - self.code = try container.decode(Int.self, forKey: .code) - self.message = try container.decode(String.self, forKey: .message) - self.status = try container.decodeIfPresent(RPCErrorStatus.self, forKey: .status) - self.details = try container.decodeIfPresent([Detail].self, forKey: .details) - self.explicitRetryDelay = nil - } - } - - package func encode(to encoder: any Encoder) throws { - var envelopeContainer = encoder.container(keyedBy: EnvelopeCodingKeys.self) - var statusContainer = envelopeContainer.nestedContainer( - keyedBy: StatusCodingKeys.self, forKey: .error) - try statusContainer.encode(code, forKey: .code) - try statusContainer.encode(message, forKey: .message) - try statusContainer.encodeIfPresent(status, forKey: .status) - try statusContainer.encodeIfPresent(details, forKey: .details) - } - /// Returns a copy of this error with the specified retry delay duration. /// /// - Parameter retryDelay: The retry delay duration. diff --git a/Tests/GeminiAPIClientTests/GeminiAPIClientIntegrationTests.swift b/Tests/GeminiAPIClientTests/GeminiAPIClientIntegrationTests.swift new file mode 100644 index 00000000000..11da2b83bde --- /dev/null +++ b/Tests/GeminiAPIClientTests/GeminiAPIClientIntegrationTests.swift @@ -0,0 +1,268 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import GenerateContentDataModels +import SharedDataModels +import SharedTestUtilities +import Testing + +@testable import GeminiAPIClient + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +@Suite("GeminiAPIClient Integration Tests") +struct GeminiAPIClientIntegrationTests { + private static let defaultBaseURL = URL(string: "https://generativelanguage.googleapis.com")! + private static let defaultModelResourcePath = "v1beta/models/gemini-3.5-flash-lite" + + @Test(.requireAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentSimplePrompt() async throws { + let client = makeClient() + let request = GenerateContentRequest( + contents: [ + Content( + parts: [Part(data: .text("Reply with the single word 'HELLO'."))], + role: "user" + ) + ] + ) + + let stream = try await client.generateContentStream(for: request) + var accumulatedText = "" + var chunkCount = 0 + for try await chunk in stream { + chunkCount += 1 + if let text = extractText(from: chunk) { + accumulatedText += text + } + } + + #expect(chunkCount > 0) + #expect(!accumulatedText.isEmpty) + #expect(accumulatedText.localizedCaseInsensitiveContains("HELLO")) + } + + @Test(.requireAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentMultiTurn() async throws { + let client = makeClient() + let contents = [ + Content(parts: [Part(data: .text("My favorite color is teal."))], role: "user"), + Content(parts: [Part(data: .text("Got it! Your favorite color is teal."))], role: "model"), + Content( + parts: [Part(data: .text("What is my favorite color? Answer in one word."))], + role: "user" + ), + ] + let request = GenerateContentRequest(contents: contents) + + let stream = try await client.generateContentStream(for: request) + var accumulatedText = "" + var chunkCount = 0 + for try await chunk in stream { + chunkCount += 1 + if let text = extractText(from: chunk) { + accumulatedText += text + } + } + + #expect(chunkCount > 0) + #expect(!accumulatedText.isEmpty) + #expect(accumulatedText.localizedCaseInsensitiveContains("teal")) + } + + @Test(.requireAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentWithSystemInstruction() async throws { + let client = makeClient() + let systemInstruction = Content( + parts: [Part(data: .text("Always speak like a 17th-century pirate."))] + ) + let request = GenerateContentRequest( + model: nil, + systemInstruction: systemInstruction, + contents: [ + Content( + parts: [Part(data: .text("How is the weather today?"))], + role: "user" + ) + ] + ) + + let stream = try await client.generateContentStream(for: request) + var accumulatedText = "" + var chunkCount = 0 + for try await chunk in stream { + chunkCount += 1 + if let text = extractText(from: chunk) { + accumulatedText += text + } + } + + #expect(chunkCount > 0) + #expect(!accumulatedText.isEmpty) + } + + @Test(.requireAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentInvalidModelNameThrows() async throws { + let client = makeClient(modelResourcePath: "v1beta/models/non-existent-model-name-xyz-123") + let request = GenerateContentRequest( + contents: [ + Content( + parts: [Part(data: .text("Hello"))], + role: "user" + ) + ] + ) + + await #expect(throws: (any Error).self) { + let stream = try await client.generateContentStream(for: request) + for try await _ in stream {} + } + } + + @Test(.requireNoAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentWithoutAuthThrows() async throws { + let client = GeminiAPIClient( + modelResourcePath: Self.defaultModelResourcePath, + baseURL: Self.defaultBaseURL + ) + let request = GenerateContentRequest( + contents: [ + Content( + parts: [Part(data: .text("Hello"))], + role: "user" + ) + ] + ) + + do { + let stream = try await client.generateContentStream(for: request) + for try await _ in stream {} + Issue.record("Expected request to throw due to missing authentication") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 400 || apiError.code == 403 || apiError.code == 401) + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test(.requireAPIKey) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func countTokensSimplePrompt() async throws { + let client = makeClient() + let request = CountTokensRequest( + contents: [ + Content( + parts: [Part(data: .text("The quick brown fox jumps over the lazy dog."))], + role: "user" + ) + ] + ) + + let response = try await client.countTokens(for: request) + + let totalTokens = try #require(response.totalTokens) + #expect(totalTokens > 0) + } + + @Test(.requireFirebaseAILogic) + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentFirebaseAILogicDeveloperAPI() async throws { + let projectID = try #require(firebaseProjectID) + let appID = try #require(firebaseAppID) + let apiKey = try #require(firebaseAPIKey) + let debugToken = try #require(appCheckDebugToken) + let appCheckClient = AppCheckDebugClient( + projectID: projectID, + appID: appID, + apiKey: apiKey, + debugToken: debugToken + ) + let appCheckToken = try await appCheckClient.exchangeDebugToken() + let modelPath = "v1beta/projects/\(projectID)/models/gemini-3.5-flash-lite" + let client = GeminiAPIClient( + modelResourcePath: modelPath, + baseURL: URL(string: "https://firebasevertexai.googleapis.com")!, + headerProvider: { + [ + "x-goog-api-key": apiKey, + "x-firebase-appcheck": appCheckToken, + ] + } + ) + let request = GenerateContentRequest( + contents: [ + Content( + parts: [Part(data: .text("Reply with the single word 'HELLO'."))], + role: "user" + ) + ] + ) + + let stream = try await client.generateContentStream(for: request) + var accumulatedText = "" + var chunkCount = 0 + for try await chunk in stream { + chunkCount += 1 + if let text = extractText(from: chunk) { + accumulatedText += text + } + } + + #expect(chunkCount > 0) + #expect(!accumulatedText.isEmpty) + #expect(accumulatedText.localizedCaseInsensitiveContains("HELLO")) + } + + // MARK: - Test Helpers + + private func extractText(from response: GenerateContentResponse?) -> String? { + guard let parts = response?.candidates?.first?.content?.parts else { return nil } + let textParts = parts.compactMap { part -> String? in + if case .text(let text) = part.data { return text } + return nil + } + return textParts.isEmpty ? nil : textParts.joined() + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + private func makeClient( + modelResourcePath: String = defaultModelResourcePath, + baseURL: URL = defaultBaseURL, + configuration: URLSessionConfiguration = .ephemeral + ) -> GeminiAPIClient { + let headerProvider: (@Sendable () async throws -> [String: String])? + if let apiKey = geminiAPIKey { + headerProvider = { @Sendable in + ["x-goog-api-key": apiKey] + } + } else { + headerProvider = nil + } + + return GeminiAPIClient( + modelResourcePath: modelResourcePath, + baseURL: baseURL, + headerProvider: headerProvider, + sessionConfiguration: configuration + ) + } +} diff --git a/Tests/GeminiAPIClientTests/GeminiAPIClientTests.swift b/Tests/GeminiAPIClientTests/GeminiAPIClientTests.swift new file mode 100644 index 00000000000..a8799b94bb2 --- /dev/null +++ b/Tests/GeminiAPIClientTests/GeminiAPIClientTests.swift @@ -0,0 +1,937 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import GenerateContentDataModels +import HTTPStreamingClient +import SharedDataModels +import SharedTestUtilities +import Testing + +@testable import GeminiAPIClient + +#if canImport(FoundationNetworking) + import FoundationNetworking +#endif + +@Suite("GeminiAPIClient Tests") +struct GeminiAPIClientTests { + private static let baseURLString = "https://generativelanguage.googleapis.com" + private static let defaultModelResourcePath = "v1beta/models/gemini-3.5-flash-lite" + + private let testID = UUID().uuidString + + private var testBaseURL: URL { + URL(string: "\(Self.baseURLString)/\(testID)")! + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentSingleChunk() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let ssePayload = """ + data: {"candidates": [{"content": {"parts": [{"text": "Hello world!"}], "role": "model"}, "finishReason": "STOP", "index": 0}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in + #expect(request.value(forHTTPHeaderField: "Content-Type") == "application/json") + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Say hello") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.content?.parts?.first?.data == .text("Hello world!")) + #expect(candidate.finishReason == .stop) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentSafetyBlockResponse() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let ssePayload = """ + data: {"candidates": [{"content": {}, "finishReason": "SAFETY", "index": 0, "finishMessage": "The model output could not be generated due to safety policy.", "safetyRatings": [{"category": "HARM_CATEGORY_DANGEROUS_CONTENT", "probability": "HIGH"}]}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Safety test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.finishReason == .safety) + #expect(candidate.finishMessage?.contains("safety policy") == true) + #expect(candidate.safetyRatings?.first?.category == .dangerousContent) + #expect(candidate.content?.parts?.first?.data == nil) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentRecitationBlockResponse() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let ssePayload = """ + data: {"candidates": [{"finishReason": "RECITATION", "index": 0}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Recitation test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.finishReason == .recitation) + #expect(candidate.content?.parts?.first?.data == nil) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentMultipleChunks() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let ssePayload = """ + data: {"candidates": [{"content": {"parts": [{"text": "Hello"}], "role": "model"}}]} + + data: {"candidates": [{"content": {"parts": [{"text": " world"}], "role": "model"}}]} + + data: {"candidates": [{"content": {"parts": [{"text": "!"}], "role": "model"}, "finishReason": "STOP"}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Say hello world") + let stream = try await client.generateContentStream(for: request) + var collectedText = "" + for try await chunk in stream { + if let text = extractText(from: chunk) { + collectedText += text + } + } + + #expect(collectedText == "Hello world!") + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentAPIError() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 400, + headerFields: ["Content-Type": "application/json"] + ) + + let errorJSON = """ + { + "error": { + "code": 400, + "message": "API key not valid. Please pass a valid API key.", + "status": "INVALID_ARGUMENT" + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(errorJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Invalid key test") + + await #expect(throws: GeminiAPIError.self) { + try await client.generateContentStream(for: request) + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentHTTPError() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 500, + headerFields: nil + ) + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data("Internal Server Error".utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Internal error test") + + do { + _ = try await client.generateContentStream(for: request) + Issue.record("Expected GeminiAPIError.httpError to be thrown") + } catch let GeminiAPIError.httpError(statusCode, body) { + #expect(statusCode == 500) + #expect(body == "Internal Server Error") + } catch { + Issue.record("Unexpected error: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentRateLimitError() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 429, + headerFields: [ + "Content-Type": "application/json", + "Retry-After": "60", + ] + ) + + let errorJSON = """ + { + "error": { + "code": 429, + "message": "Resource has been exhausted (e.g. check quota).", + "status": "RESOURCE_EXHAUSTED" + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(errorJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Rate limit test") + + do { + _ = try await client.generateContentStream(for: request) + Issue.record("Expected GeminiAPIError.apiError to be thrown") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 429) + #expect(apiError.status == .resourceExhausted) + #expect(apiError.message.contains("Resource has been exhausted")) + #expect(apiError.retryDelay == .seconds(60)) + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentRateLimitWithRetryInfoJSON() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 429, + headerFields: ["Content-Type": "application/json"] + ) + + let errorJSON = """ + { + "error": { + "code": 429, + "message": "Resource has been exhausted (e.g. check quota).", + "status": "RESOURCE_EXHAUSTED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay": "45.5s" + } + ] + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(errorJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Rate limit test") + + do { + _ = try await client.generateContentStream(for: request) + Issue.record("Expected GeminiAPIError.apiError to be thrown") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 429) + #expect(apiError.status == .resourceExhausted) + #expect(apiError.retryDelay == .seconds(45.5)) + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func headerProviderInjectsHeaders() async throws { + let client = makeClient( + headerProvider: { + [ + "x-goog-api-key": "custom-key", + "X-AppCheck-Token": "app-check-123", + ] + } + ) + + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 200, + headerFields: nil + ) + + MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in + #expect(request.value(forHTTPHeaderField: "x-goog-api-key") == "custom-key") + #expect(request.value(forHTTPHeaderField: "X-AppCheck-Token") == "app-check-123") + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol( + proto, + didLoad: Data( + "data: {\"candidates\": [{\"content\": {\"parts\": [{\"text\": \"OK\"}]}}]}\n\n".utf8) + ) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let stream = try await client.generateContentStream( + for: makePromptRequest("Header test")) + var count = 0 + for try await chunk in stream { + count += 1 + #expect(extractText(from: chunk) == "OK") + } + + #expect(count == 1) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentClientTimeoutThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didFailWithError: URLError(.timedOut)) + } + + let request = makePromptRequest("Timeout test") + + await #expect(throws: URLError.self) { + try await client.generateContentStream(for: request) + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentGatewayTimeoutThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 504, + headerFields: ["Content-Type": "application/json"] + ) + + let errorJSON = """ + { + "error": { + "code": 504, + "message": "Deadline exceeded while waiting for model response.", + "status": "DEADLINE_EXCEEDED" + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(errorJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Gateway timeout test") + + do { + _ = try await client.generateContentStream(for: request) + Issue.record("Expected GeminiAPIError.apiError to be thrown") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 504) + #expect(apiError.status == .deadlineExceeded) + #expect(apiError.message.contains("Deadline exceeded")) + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentMidStreamErrorThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let payload = """ + data: {"candidates": [{"content": {"parts": [{"text": "First "}]},"finishReason": "STOP","index": 0}]} + + data: {"candidates": [{"content": {"parts": [{"text": "Second "}]},"finishReason": "STOP","index": 0}]} + + { + "error": { + "code": 499, + "message": "The operation was cancelled.", + "status": "CANCELLED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.DebugInfo", + "detail": "[ORIGINAL ERROR] generic::cancelled: " + } + ] + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Mid-stream error test") + let stream = try await client.generateContentStream(for: request) + var collectedText = "" + + do { + for try await chunk in stream { + if let text = extractText(from: chunk) { + collectedText += text + } + } + Issue.record("Expected GeminiAPIError.apiError to be thrown mid-stream") + } catch let GeminiAPIError.apiError(apiError) { + #expect(collectedText == "First Second ") + #expect(apiError.code == 499) + #expect(apiError.status == .cancelled) + #expect(apiError.message == "The operation was cancelled.") + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentErrorInSSEDataThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let payload = """ + data: {"error": {"code": 500, "message": "Internal error occurred.", "status": "INTERNAL"}} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("SSE error test") + let stream = try await client.generateContentStream(for: request) + + do { + for try await _ in stream {} + Issue.record("Expected GeminiAPIError.apiError to be thrown") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 500) + #expect(apiError.status == .internalError) + #expect(apiError.message == "Internal error occurred.") + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentMidStreamUnrecognizedTextThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let payload = """ + data: {"candidates": [{"content": {"parts": [{"text": "Hello"}]}}]} + + Unrecognized non-JSON error payload + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(payload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Unrecognized payload test") + let stream = try await client.generateContentStream(for: request) + var collectedText = "" + + do { + for try await chunk in stream { + if let text = extractText(from: chunk) { + collectedText += text + } + } + Issue.record("Expected GeminiAPIError.httpError to be thrown") + } catch let GeminiAPIError.httpError(statusCode, body) { + #expect(collectedText == "Hello") + #expect(statusCode == 200) + #expect(body == "Unrecognized non-JSON error payload") + } catch { + Issue.record("Unexpected error thrown: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func countTokensSuccess() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL(action: "countTokens", query: nil) + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "application/json"] + ) + + let jsonPayload = """ + { + "totalTokens": 42, + "cachedContentTokenCount": 10 + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in + #expect(request.value(forHTTPHeaderField: "Content-Type") == "application/json") + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(jsonPayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = CountTokensRequest( + contents: [Content(parts: [Part(data: .text("Count me"))], role: "user")] + ) + let response = try await client.countTokens(for: request) + + #expect(response.totalTokens == 42) + #expect(response.cachedContentTokenCount == 10) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func countTokensAPIError() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL(action: "countTokens", query: nil) + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 400, + headerFields: ["Content-Type": "application/json"] + ) + + let errorJSON = """ + { + "error": { + "code": 400, + "message": "Invalid argument for token count", + "status": "INVALID_ARGUMENT" + } + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(errorJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = CountTokensRequest(contents: []) + + do { + _ = try await client.countTokens(for: request) + Issue.record("Expected GeminiAPIError.apiError to be thrown") + } catch let GeminiAPIError.apiError(apiError) { + #expect(apiError.code == 400) + #expect(apiError.status == .invalidArgument) + #expect(apiError.message == "Invalid argument for token count") + } catch { + Issue.record("Unexpected error: \(error)") + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func countTokensDecodingErrorThrows() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL(action: "countTokens", query: nil) + let httpResponse = try makeResponse( + url: expectedURL, + statusCode: 200, + headerFields: ["Content-Type": "application/json"] + ) + + let invalidJSON = "{\"totalTokens\": \"invalid-type-string\"}" + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(invalidJSON.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = CountTokensRequest(contents: []) + + await #expect(throws: DecodingError.self) { + _ = try await client.countTokens(for: request) + } + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentMultiLineSSEData() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + + let ssePayload = """ + data: {"candidates": [{"content": {"parts": + data: [{"text": "Multi-line SSE"}], "role": "model"}, "finishReason": "STOP", "index": 0}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Multi-line test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.content?.parts?.first?.data == .text("Multi-line SSE")) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentWithComments() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + let ssePayload = """ + : keep-alive ping + + data: {"candidates": [{"content": {"parts": [{"text": "With comments"}], "role": "model"}, "finishReason": "STOP", "index": 0}]} + + : end of stream + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Comments test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.content?.parts?.first?.data == .text("With comments")) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentWithSSEControlFields() async throws { + let client = makeClient() + let expectedURL = try makeExpectedURL() + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + let ssePayload = """ + event: message + id: 101 + retry: 3000 + data: {"candidates": [{"content": {"parts": [{"text": "SSE control fields ignored"}], "role": "model"}, "finishReason": "STOP", "index": 0}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Control fields test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let candidate = try #require(responses.first?.candidates?.first) + #expect(candidate.content?.parts?.first?.data == .text("SSE control fields ignored")) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentAgentPlatformResourcePath() async throws { + let agentPlatformPath = + "v1beta1/projects/my-project/locations/global/publishers/google/models/gemini-3.5-flash-lite" + let client = makeClient(modelResourcePath: agentPlatformPath) + let expectedURL = try makeExpectedURL(modelResourcePath: agentPlatformPath) + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + let ssePayload = """ + data: {"candidates": [{"content": {"parts": [{"text": "Agent Platform response"}]}}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Agent Platform test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let response = try #require(responses.first) + let candidate = try #require(response.candidates?.first) + let content = try #require(candidate.content) + let part = try #require(content.parts?.first) + #expect(part.data == .text("Agent Platform response")) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func streamGenerateContentFirebaseProxyResourcePath() async throws { + let firebaseProxyPath = "v1beta/projects/my-firebase-project/models/gemini-3.5-flash-lite" + let client = makeClient(modelResourcePath: firebaseProxyPath) + let expectedURL = try makeExpectedURL(modelResourcePath: firebaseProxyPath) + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "text/event-stream"] + ) + let ssePayload = """ + data: {"candidates": [{"content": {"parts": [{"text": "Firebase AI Logic response"}]}}]} + + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { _, proto in + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(ssePayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = makePromptRequest("Firebase AI Logic test") + let stream = try await client.generateContentStream(for: request) + var responses: [GenerateContentResponse] = [] + for try await chunk in stream { + responses.append(chunk) + } + + #expect(responses.count == 1) + let response = try #require(responses.first) + let candidate = try #require(response.candidates?.first) + let content = try #require(candidate.content) + let part = try #require(content.parts?.first) + #expect(part.data == .text("Firebase AI Logic response")) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func countTokensAgentPlatformResourcePath() async throws { + let agentPlatformPath = + "v1beta1/projects/my-project/locations/global/publishers/google/models/gemini-3.5-flash-lite" + let client = makeClient(modelResourcePath: agentPlatformPath) + let expectedURL = try makeExpectedURL( + modelResourcePath: agentPlatformPath, action: "countTokens", query: nil + ) + let httpResponse = try makeResponse( + url: expectedURL, + headerFields: ["Content-Type": "application/json"] + ) + let jsonPayload = """ + { + "totalTokens": 128 + } + """ + + MockHTTPURLProtocol.setHandler(for: expectedURL) { request, proto in + #expect(request.value(forHTTPHeaderField: "Content-Type") == "application/json") + proto.client?.urlProtocol(proto, didReceive: httpResponse, cacheStoragePolicy: .notAllowed) + proto.client?.urlProtocol(proto, didLoad: Data(jsonPayload.utf8)) + proto.client?.urlProtocolDidFinishLoading(proto) + } + + let request = CountTokensRequest( + contents: [Content(parts: [Part(data: .text("Agent Platform count tokens"))], role: "user")] + ) + let response = try await client.countTokens(for: request) + + #expect(response.totalTokens == 128) + } + + // MARK: - Test Helpers + + private func makePromptRequest(_ prompt: String) -> GenerateContentRequest { + GenerateContentRequest( + contents: [Content(parts: [Part(data: .text(prompt))], role: "user")] + ) + } + + private func extractText(from response: GenerateContentResponse?) -> String? { + guard let parts = response?.candidates?.first?.content?.parts else { return nil } + let textParts = parts.compactMap { part -> String? in + if case .text(let text) = part.data { return text } + return nil + } + return textParts.isEmpty ? nil : textParts.joined() + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + private func makeClient( + modelResourcePath: String = defaultModelResourcePath, + baseURL: URL? = nil, + headerProvider: (@Sendable () async throws -> [String: String])? = nil + ) -> GeminiAPIClient { + let configuration = URLSessionConfiguration.ephemeral + configuration.protocolClasses = [MockHTTPURLProtocol.self] + return GeminiAPIClient( + modelResourcePath: modelResourcePath, + baseURL: baseURL ?? testBaseURL, + headerProvider: headerProvider, + sessionConfiguration: configuration + ) + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + private func makeResponse( + url: URL, + statusCode: Int = 200, + headerFields: [String: String]? = nil + ) throws -> HTTPURLResponse { + try HTTPURLResponse.mock(url: url, statusCode: statusCode, headerFields: headerFields) + } + + private func makeExpectedURL( + modelResourcePath: String = defaultModelResourcePath, + action: String = "streamGenerateContent", + query: String? = "alt=sse" + ) throws -> URL { + var urlString = "\(testBaseURL.absoluteString)/\(modelResourcePath):\(action)" + if let query { + urlString += "?\(query)" + } + return try #require(URL(string: urlString)) + } +} diff --git a/Tests/GeminiAPIClientTests/GeminiAPIErrorTests.swift b/Tests/GeminiAPIClientTests/GeminiAPIErrorTests.swift new file mode 100644 index 00000000000..72790fe9244 --- /dev/null +++ b/Tests/GeminiAPIClientTests/GeminiAPIErrorTests.swift @@ -0,0 +1,108 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +import Foundation +import SharedDataModels +import Testing + +@testable import GeminiAPIClient + +@Suite("GeminiAPIError Tests") +struct GeminiAPIErrorTests { + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func retryInfoDecodingAndDurationCalculation() throws { + let json = """ + { + "error": { + "code": 429, + "message": "Resource exhausted", + "status": "RESOURCE_EXHAUSTED", + "details": [ + { + "@type": "type.googleapis.com/google.rpc.RetryInfo", + "retryDelay": "12.5s" + } + ] + } + } + """ + + let error = try JSONDecoder().decode(GoogleCloudAPIError.self, from: Data(json.utf8)) + + #expect(error.retryDelay == .seconds(12.5)) + #expect(error.code == 429) + #expect(error.status == .resourceExhausted) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func geminiAPIErrorRetryAfterPrecedence() { + let cloudErrorWithDelay = GoogleCloudAPIError( + code: 429, + message: "Quota exceeded", + status: .resourceExhausted, + details: [.retryInfo(GoogleCloudAPIError.RetryInfo(retryDelay: "30s"))] + ) + + let apiErrorWithOverride = GeminiAPIError.apiError( + cloudErrorWithDelay.withRetryDelay(.seconds(60))) + let apiErrorWithoutOverride = GeminiAPIError.apiError(cloudErrorWithDelay) + let httpError = GeminiAPIError.httpError(statusCode: 500, body: "Server error") + + #expect(apiErrorWithOverride.retryAfter == .seconds(60)) + #expect(apiErrorWithoutOverride.retryAfter == .seconds(30)) + #expect(httpError.retryAfter == nil) + } + + @Test + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + func localizedErrorAndCustomNSErrorProperties() { + let cloudError = GoogleCloudAPIError( + code: 400, + message: "Invalid argument message", + status: .invalidArgument, + details: [ + .localizedMessage( + GoogleCloudAPIError.LocalizedMessage( + locale: "en-US", message: "Localized argument error")), + .help( + GoogleCloudAPIError.Help(links: [ + GoogleCloudAPIError.Help.Link( + description: "Help doc", url: "https://cloud.google.com/help") + ])), + ], + retryDelay: .seconds(15) + ) + + let apiError = GeminiAPIError.apiError(cloudError) + let httpError = GeminiAPIError.httpError(statusCode: 503, body: "Service unavailable") + + #expect(apiError.errorDescription == "Localized argument error") + #expect(apiError.failureReason == "INVALID_ARGUMENT") + #expect(apiError.helpAnchor == "https://cloud.google.com/help") + #expect(apiError.errorCode == 400) + #expect(GeminiAPIError.errorDomain == "GeminiAPIClient.GeminiAPIError") + #expect(apiError.errorUserInfo["code"] as? Int == 400) + #expect(apiError.errorUserInfo["status"] as? String == "INVALID_ARGUMENT") + #expect(apiError.errorUserInfo["retryAfterSeconds"] as? Double == 15.0) + + #expect(httpError.errorDescription == "HTTP 503: Service unavailable") + #expect(httpError.failureReason == "HTTP status 503") + #expect(httpError.helpAnchor == nil) + #expect(httpError.errorCode == 503) + #expect(httpError.errorUserInfo["statusCode"] as? Int == 503) + #expect(httpError.errorUserInfo["body"] as? String == "Service unavailable") + } +} diff --git a/Tests/HTTPStreamingClientTests/HTTPStreamingClientTests.swift b/Tests/HTTPStreamingClientTests/HTTPStreamingClientTests.swift index e46c00f221a..c8221caddad 100644 --- a/Tests/HTTPStreamingClientTests/HTTPStreamingClientTests.swift +++ b/Tests/HTTPStreamingClientTests/HTTPStreamingClientTests.swift @@ -21,29 +21,29 @@ import Testing import FoundationNetworking #endif -// MARK: - HTTPStreamingClient Integration Tests +// MARK: - HTTPStreamingClient Unit Tests -@Suite("HTTPStreamingClient Integration Tests", .serialized) +@Suite("HTTPStreamingClient Tests") struct HTTPStreamingClientTests { + private let testID = UUID().uuidString + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) private func makeClient() -> HTTPStreamingClient { - MockHTTPURLProtocol.reset() let configuration = URLSessionConfiguration.ephemeral configuration.protocolClasses = [MockHTTPURLProtocol.self] return HTTPStreamingClient(configuration: configuration) } + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + private func makeTestURL(_ path: String) throws -> URL { + try #require(URL(string: "https://example.com/\(testID)/\(path)")) + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) private func makeResponse(url: URL, statusCode: Int = 200, headerFields: [String: String]? = nil) throws -> HTTPURLResponse { - try #require( - HTTPURLResponse( - url: url, - statusCode: statusCode, - httpVersion: "HTTP/1.1", - headerFields: headerFields - ) - ) + try HTTPURLResponse.mock(url: url, statusCode: statusCode, headerFields: headerFields) } @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) @@ -59,7 +59,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func streamsLinesAndReturnsTask() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/bytes-api")) + let testURL = try makeTestURL("bytes-api") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -79,7 +79,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func streamSingleChunkMultipleLines() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/stream-single-chunk")) + let testURL = try makeTestURL("stream-single-chunk") let testResponse = try makeResponse( url: testURL, statusCode: 200, headerFields: ["Content-Type": "text/plain"] ) @@ -101,7 +101,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func linesSplitAcrossMultipleChunks() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/stream-split-chunks")) + let testURL = try makeTestURL("stream-split-chunks") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -123,7 +123,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func streamSplitCRLFAndUTF8() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/stream-crlf-utf8")) + let testURL = try makeTestURL("stream-crlf-utf8") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) let emojiBytes: [UInt8] = [0xF0, 0x9F, 0x8E, 0x89] // 🎉 @@ -146,7 +146,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func streamPreservesBlankLines() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/stream-blank-lines")) + let testURL = try makeTestURL("stream-blank-lines") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -169,7 +169,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func streamNon2xxResponseReturnsStatusCodeAndBody() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/error-404")) + let testURL = try makeTestURL("error-404") let testResponse = try makeResponse(url: testURL, statusCode: 404, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -189,7 +189,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func networkErrorBeforeResponseThrows() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/fail-before-response")) + let testURL = try makeTestURL("fail-before-response") MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in proto.client?.urlProtocol(proto, didFailWithError: URLError(.cannotConnectToHost)) @@ -208,7 +208,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func networkErrorDuringStreamThrows() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/fail-during-stream")) + let testURL = try makeTestURL("fail-during-stream") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -232,7 +232,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func earlyTerminationCancelsStream() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/early-termination")) + let testURL = try makeTestURL("early-termination") let testResponse = try makeResponse(url: testURL, statusCode: 200, headerFields: nil) MockHTTPURLProtocol.setHandler(for: testURL) { _, proto in @@ -257,7 +257,7 @@ struct HTTPStreamingClientTests { @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) func nonHTTPResponseThrowsBadServerResponse() async throws { let client = makeClient() - let testURL = try #require(URL(string: "https://example.com/non-http")) + let testURL = try makeTestURL("non-http") let nonHTTPResponse = try #require( URLResponse( url: testURL, diff --git a/Tests/SharedTestUtilities/AppCheckDebugClient.swift b/Tests/SharedTestUtilities/AppCheckDebugClient.swift new file mode 100644 index 00000000000..f50ca7e9f79 --- /dev/null +++ b/Tests/SharedTestUtilities/AppCheckDebugClient.swift @@ -0,0 +1,133 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if canImport(Darwin) + package import Foundation +#else + import Foundation +#endif + +#if canImport(FoundationNetworking) + package import FoundationNetworking +#endif + +/// Lightweight actor client for exchanging and caching an App Check Debug Token in integration +/// tests. +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +package actor AppCheckDebugClient { + private let projectID: String + private let appID: String + private let apiKey: String + private let debugToken: String + + private var cachedToken: String? + private var inFlightExchangeTask: Task? + private static let baseURL = URL(string: "https://firebaseappcheck.googleapis.com")! + + /// Initializes the client with the required project, app, and authentication parameters. + /// + /// - Parameters: + /// - projectID: The Firebase project ID. + /// - appID: The Firebase app ID. + /// - apiKey: The Firebase API key. + /// - debugToken: The App Check debug token. + package init( + projectID: String, + appID: String, + apiKey: String, + debugToken: String + ) { + assert(!projectID.isEmpty, "projectID must not be empty.") + assert(!appID.isEmpty, "appID must not be empty.") + assert(!apiKey.isEmpty, "apiKey must not be empty.") + assert(!debugToken.isEmpty, "debugToken must not be empty.") + self.projectID = projectID + self.appID = appID + self.apiKey = apiKey + self.debugToken = debugToken + } + + /// Exchanges the debug token for an App Check token, caching the result in-memory. + /// + /// - Parameters: + /// - session: The `URLSession` to use. Defaults to `.shared`. + /// - forceRefresh: Whether to bypass the cached token and fetch a new one. + /// - Returns: The exchanged App Check token string. + /// - Throws: An error if the request fails or the response cannot be parsed. + package func exchangeDebugToken( + session: URLSession = .shared, + forceRefresh: Bool = false + ) async throws -> String { + if !forceRefresh, let cachedToken { + return cachedToken + } + if !forceRefresh, let inFlightExchangeTask { + return try await inFlightExchangeTask.value + } + + let projectID = self.projectID + let appID = self.appID + let apiKey = self.apiKey + let debugToken = self.debugToken + + let exchangeTask = Task { + let endpointURL = Self.baseURL + .appending(component: "v1") + .appending(component: "projects") + .appending(component: projectID) + .appending(component: "apps") + .appending(component: "\(appID):exchangeDebugToken") + + var request = URLRequest(url: endpointURL) + request.httpMethod = "POST" + request.setValue("application/json", forHTTPHeaderField: "Content-Type") + request.setValue(apiKey, forHTTPHeaderField: "x-goog-api-key") + request.httpBody = try JSONEncoder().encode(RequestBody(debugToken: debugToken)) + + let (data, response) = try await session.data(for: request) + guard let httpResponse = response as? HTTPURLResponse, httpResponse.statusCode == 200 else { + let errorBody = String(decoding: data, as: UTF8.self) + throw URLError(.badServerResponse, userInfo: ["body": errorBody]) + } + + let decoded = try JSONDecoder().decode(ResponseBody.self, from: data) + return decoded.token + } + + self.inFlightExchangeTask = exchangeTask + do { + let token = try await exchangeTask.value + self.cachedToken = token + self.inFlightExchangeTask = nil + return token + } catch { + self.inFlightExchangeTask = nil + throw error + } + } +} + +// MARK: - Private Payload Types + +@available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) +extension AppCheckDebugClient { + fileprivate struct RequestBody: Encodable { + let debugToken: String + } + + fileprivate struct ResponseBody: Decodable { + let token: String + let ttl: String? + } +} diff --git a/Tests/SharedTestUtilities/ConditionTrait+Credentials.swift b/Tests/SharedTestUtilities/ConditionTrait+Credentials.swift new file mode 100644 index 00000000000..c0766a2e666 --- /dev/null +++ b/Tests/SharedTestUtilities/ConditionTrait+Credentials.swift @@ -0,0 +1,112 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if canImport(Testing) + import Foundation + package import Testing + + /// Resolves the Gemini API key from `GOOGLE_API_KEY` or `GEMINI_API_KEY` environment variables. + package var geminiAPIKey: String? { + let env = ProcessInfo.processInfo.environment + if let googleKey = env["GOOGLE_API_KEY"], !googleKey.isEmpty { + return googleKey + } + if let geminiKey = env["GEMINI_API_KEY"], !geminiKey.isEmpty { + return geminiKey + } + return nil + } + + /// Indicates whether a Gemini API key is available in the environment. + package var hasGeminiAPIKey: Bool { + geminiAPIKey != nil + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + extension Trait where Self == Testing.ConditionTrait { + /// Requires a Gemini API key (`GOOGLE_API_KEY` or `GEMINI_API_KEY`) to be set in the environment. + package static var requireAPIKey: Self { + .enabled( + if: hasGeminiAPIKey, + "Requires GOOGLE_API_KEY or GEMINI_API_KEY environment variable" + ) + } + + /// Requires that no Gemini API key is set in the environment. + package static var requireNoAPIKey: Self { + .enabled( + if: !hasGeminiAPIKey, + "Runs only when no API key is set" + ) + } + } + + /// Resolves the Firebase Project ID from the `FIREBASE_PROJECT_ID` environment variable. + package var firebaseProjectID: String? { + let env = ProcessInfo.processInfo.environment + if let projectID = env["FIREBASE_PROJECT_ID"], !projectID.isEmpty { + return projectID + } + return nil + } + + /// Resolves the Firebase App ID from the `FIREBASE_APP_ID` environment variable. + package var firebaseAppID: String? { + let env = ProcessInfo.processInfo.environment + if let appID = env["FIREBASE_APP_ID"], !appID.isEmpty { + return appID + } + return nil + } + + /// Resolves the Firebase API key from the `FIREBASE_API_KEY` environment variable. + package var firebaseAPIKey: String? { + let env = ProcessInfo.processInfo.environment + if let apiKey = env["FIREBASE_API_KEY"], !apiKey.isEmpty { + return apiKey + } + return nil + } + + /// Resolves the Firebase App Check debug token from standard environment variables: + /// `AppCheckDebugToken` with fallback to `FIRAAppCheckDebugToken`. + package var appCheckDebugToken: String? { + let env = ProcessInfo.processInfo.environment + if let token = env["AppCheckDebugToken"], !token.isEmpty { + return token + } + if let legacyToken = env["FIRAAppCheckDebugToken"], !legacyToken.isEmpty { + return legacyToken + } + return nil + } + + /// Indicates whether all required Firebase AI Logic credentials and debug token are available. + package var hasFirebaseAILogicCredentials: Bool { + firebaseProjectID != nil && firebaseAppID != nil && firebaseAPIKey != nil + && appCheckDebugToken != nil + } + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + extension Trait where Self == Testing.ConditionTrait { + /// Requires all Firebase AI Logic credentials and an App Check debug token to be set in the + /// environment variables. + package static var requireFirebaseAILogic: Self { + .enabled( + if: hasFirebaseAILogicCredentials, + "Requires FIREBASE_PROJECT_ID, FIREBASE_APP_ID, FIREBASE_API_KEY, and AppCheckDebugToken." + ) + } + } +#endif // canImport(Testing) diff --git a/Tests/SharedTestUtilities/HTTPURLResponse+Mock.swift b/Tests/SharedTestUtilities/HTTPURLResponse+Mock.swift new file mode 100644 index 00000000000..294f68b1148 --- /dev/null +++ b/Tests/SharedTestUtilities/HTTPURLResponse+Mock.swift @@ -0,0 +1,48 @@ +// Copyright 2026 Google LLC +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +#if canImport(Testing) + package import Foundation + import Testing + + #if canImport(FoundationNetworking) + package import FoundationNetworking + #endif + + @available(macOS 15.0, iOS 18.0, tvOS 18.0, watchOS 11.0, visionOS 2.0, *) + extension HTTPURLResponse { + /// Creates a mock `HTTPURLResponse` for testing. + /// + /// - Parameters: + /// - url: The target response URL. + /// - statusCode: The HTTP status code. Defaults to `200`. + /// - headerFields: The optional HTTP header dictionary. + /// - Returns: A non-nil `HTTPURLResponse`. + /// - Throws: An error if response construction fails. + package static func mock( + url: URL, + statusCode: Int = 200, + headerFields: [String: String]? = nil + ) throws -> HTTPURLResponse { + try #require( + HTTPURLResponse( + url: url, + statusCode: statusCode, + httpVersion: "HTTP/1.1", + headerFields: headerFields + ) + ) + } + } +#endif // canImport(Testing) diff --git a/Tests/SharedTestUtilities/MockHTTPURLProtocol.swift b/Tests/SharedTestUtilities/MockHTTPURLProtocol.swift index cb157a84a21..d43e65aa9e3 100644 --- a/Tests/SharedTestUtilities/MockHTTPURLProtocol.swift +++ b/Tests/SharedTestUtilities/MockHTTPURLProtocol.swift @@ -55,9 +55,10 @@ package final class MockHTTPURLProtocol: URLProtocol { /// Determines whether this protocol can handle the given request. /// /// - Parameter request: The proposed request. - /// - Returns: Always `true` for all intercepted requests. + /// - Returns: `true` if a handler is registered for this request's URL, otherwise `false`. override package class func canInit(with request: URLRequest) -> Bool { - true + guard let urlString = request.url?.absoluteString else { return false } + return handlers.withLock { $0[urlString] != nil } } /// Returns the canonical version of the given request.