[Firebase AI] Add IgnoreBackendOverloadedTrait in integration tests

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,

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)

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

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)

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
+  }
+}

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 = "<group>"; };
 		862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IntegrationTestUtils.swift; sourceTree = "<group>"; };
 		864F8F702D4980D60002EA7E /* ImagenIntegrationTests.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ImagenIntegrationTests.swift; sourceTree = "<group>"; };
 		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 */,