diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift index 9be7f1dcb02..97ca2093fde 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/CountTokensIntegrationTests.swift @@ -21,7 +21,7 @@ import VertexAITestApp @testable import struct FirebaseAI.APIConfig -@Suite(.serialized) +@Suite(.serialized, .ignoreBackendOverloaded) struct CountTokensIntegrationTests { let generationConfig = GenerationConfig( temperature: 1.2, diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift index c07803a8381..eb7a834b37d 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/GenerateContentIntegrationTests.swift @@ -25,7 +25,7 @@ import VertexAITestApp @testable import struct FirebaseAI.BackendError -@Suite(.serialized) +@Suite(.serialized, .ignoreBackendOverloaded) struct GenerateContentIntegrationTests { // Set temperature, topP and topK to lowest allowed values to make responses more deterministic. let generationConfig = GenerationConfig(temperature: 0.0, topP: 0.0, topK: 1) @@ -116,8 +116,8 @@ struct GenerateContentIntegrationTests { @Test(arguments: [ // TODO(andrewheard): Vertex AI configs temporarily disabled to due empty SafetyRatings bug. - // InstanceConfig.vertexV1, - // InstanceConfig.vertexV1Beta, + // InstanceConfig.vertexAI_v1, + // InstanceConfig.vertexAI_v1beta, InstanceConfig.googleAI_v1beta, InstanceConfig.googleAI_v1beta_staging, InstanceConfig.googleAI_v1beta_freeTier_bypassProxy, @@ -136,17 +136,8 @@ struct GenerateContentIntegrationTests { ) let prompt = "Generate an image of a cute cartoon kitten playing with a ball of yarn." - var response: GenerateContentResponse? - try await withKnownIssue( - "Backend may fail with a 503 - Service Unavailable error when overloaded", - isIntermittent: true - ) { - response = try await model.generateContent(prompt) - } matching: { issue in - (issue.error as? BackendError).map { $0.httpResponseCode == 503 } ?? false - } + let response = try await model.generateContent(prompt) - guard let response else { return } let candidate = try #require(response.candidates.first) let inlineDataPart = try #require(candidate.content.parts .first { $0 is InlineDataPart } as? InlineDataPart) diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift index b7533c08aa1..e994e143c9f 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/ImagenIntegrationTests.swift @@ -31,7 +31,8 @@ import VertexAITestApp if: ProcessInfo.processInfo.environment["VTXIntegrationImagen"] != nil, "Only runs if the environment variable VTXIntegrationImagen is set." ), - .serialized + .serialized, + .ignoreBackendOverloaded ) struct ImagenIntegrationTests { var vertex: FirebaseAI diff --git a/FirebaseAI/Tests/TestApp/Tests/Integration/SchemaTests.swift b/FirebaseAI/Tests/TestApp/Tests/Integration/SchemaTests.swift index 431134d315a..f6309fab293 100644 --- a/FirebaseAI/Tests/TestApp/Tests/Integration/SchemaTests.swift +++ b/FirebaseAI/Tests/TestApp/Tests/Integration/SchemaTests.swift @@ -26,7 +26,7 @@ import VertexAITestApp @testable import struct FirebaseAI.BackendError /// Test the schema fields. -@Suite(.serialized) +@Suite(.serialized, .ignoreBackendOverloaded) struct SchemaTests { // Set temperature, topP and topK to lowest allowed values to make responses more deterministic. let generationConfig = GenerationConfig(temperature: 0.0, topP: 0.0, topK: 1) diff --git a/FirebaseAI/Tests/TestApp/Tests/Utilities/IgnoreBackendOverloadedTrait.swift b/FirebaseAI/Tests/TestApp/Tests/Utilities/IgnoreBackendOverloadedTrait.swift new file mode 100644 index 00000000000..417de9039c1 --- /dev/null +++ b/FirebaseAI/Tests/TestApp/Tests/Utilities/IgnoreBackendOverloadedTrait.swift @@ -0,0 +1,52 @@ +// Copyright 2025 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 Testing + +@testable import FirebaseAI + +/// A trait that ignores HTTP 503 - Service Unavailable errors. +/// +/// This occurs when the backend is overloaded and cannot handle requests. +struct IgnoreBackendOverloadedTrait: TestTrait, SuiteTrait, TestScoping { + func provideScope(for test: Test, testCase: Test.Case?, + performing function: @Sendable () async throws -> Void) async throws { + try await withKnownIssue( + "Backend may fail with a 503 - Service Unavailable error when overloaded", + isIntermittent: true + ) { + try await function() + } matching: { issue in + if case let .internalError(error as BackendError) = issue.error as? GenerateContentError, + error.isServiceUnavailable { + return true + } else if let error = issue.error as? BackendError, error.isServiceUnavailable { + return true + } + + return false + } + } +} + +extension Trait where Self == IgnoreBackendOverloadedTrait { + static var ignoreBackendOverloaded: Self { Self() } +} + +extension BackendError { + /// Returns true when the error is HTTP 503 - Service Unavailable. + var isServiceUnavailable: Bool { + httpResponseCode == 503 + } +} diff --git a/FirebaseAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj b/FirebaseAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj index 50303ad511b..24ba781e605 100644 --- a/FirebaseAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj +++ b/FirebaseAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj @@ -7,6 +7,7 @@ objects = { /* Begin PBXBuildFile section */ + 8603266C2DC5671C00C2BCFA /* IgnoreBackendOverloadedTrait.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8603266B2DC5671C00C2BCFA /* IgnoreBackendOverloadedTrait.swift */; }; 862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */; }; 864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 864F8F702D4980D60002EA7E /* ImagenIntegrationTests.swift */; }; 8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8661385B2CC943DD00F4B78E /* TestApp.swift */; }; @@ -42,6 +43,7 @@ /* End PBXContainerItemProxy section */ /* Begin PBXFileReference section */ + 8603266B2DC5671C00C2BCFA /* IgnoreBackendOverloadedTrait.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IgnoreBackendOverloadedTrait.swift; sourceTree = ""; }; 862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IntegrationTestUtils.swift; sourceTree = ""; }; 864F8F702D4980D60002EA7E /* ImagenIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImagenIntegrationTests.swift; sourceTree = ""; }; 866138582CC943DD00F4B78E /* VertexAITestApp-SPM.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "VertexAITestApp-SPM.app"; sourceTree = BUILT_PRODUCTS_DIR; }; @@ -164,6 +166,7 @@ 8698D7442CD3CEF700ABA833 /* Utilities */ = { isa = PBXGroup; children = ( + 8603266B2DC5671C00C2BCFA /* IgnoreBackendOverloadedTrait.swift */, 86D77E032D7B6C95003D155D /* InstanceConfig.swift */, 862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */, ); @@ -297,6 +300,7 @@ 86D77E042D7B6C9D003D155D /* InstanceConfig.swift in Sources */, DEF0BB512DA9B7450093E9F4 /* SchemaTests.swift in Sources */, DEF0BB4F2DA74F680093E9F4 /* TestHelpers.swift in Sources */, + 8603266C2DC5671C00C2BCFA /* IgnoreBackendOverloadedTrait.swift in Sources */, 868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */, 864F8F712D4980DD0002EA7E /* ImagenIntegrationTests.swift in Sources */, 862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */,