[Vertex AI] Run Imagen integration tests on cron schedule only (#14231)

Author: andrewheard

.github/workflows/vertexai.yml

@@ -9,6 +9,7 @@ on:
   schedule:
     # Run every day at 11pm (PST) - cron uses UTC times
     - cron:  '0 7 * * *'
+  workflow_dispatch:
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.head_ref || github.ref }}
@@ -102,6 +103,7 @@ jobs:
     needs: spm-package-resolved
     env:
       TEST_RUNNER_FIRAAppCheckDebugToken: ${{ secrets.VERTEXAI_INTEGRATION_FAC_DEBUG_TOKEN }}
+      TEST_RUNNER_VTXIntegrationImagen: ${{ github.event_name == 'schedule' || github.event_name == 'workflow_dispatch' }}
       FIREBASECI_USE_LATEST_GOOGLEAPPMEASUREMENT: 1
       secrets_passphrase: ${{ secrets.GHASecretsGPGPassphrase1 }}
     steps:

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTestUtils.swift

@@ -0,0 +1,39 @@
+// Copyright 2024 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 XCTest
+
+enum IntegrationTestUtils {
+  /// Skips an XCTest unless the specified environment variable is set.
+  ///
+  /// - Parameters:
+  ///   - environmentVariable: The environment variable that must be defined for the test to
+  ///     continue (i.e., not get skipped).
+  ///   - requiredValue: If specified, skips the test if `environmentVariable` is not set to the
+  ///     this value; if `nil`, any value allows the test to continue.
+  /// - Throws: `XCTSkip` if the test should be skipped.
+  static func skipUnless(environmentVariable: String, requiredValue: String? = nil) throws {
+    guard let variableValue = ProcessInfo.processInfo.environment[environmentVariable] else {
+      throw XCTSkip("Skipped because environment variable '\(environmentVariable)' is not defined.")
+    }
+
+    if let requiredValue, variableValue != requiredValue {
+      throw XCTSkip("""
+      Skipped because environment variable '\(environmentVariable)' != '\(requiredValue)'; value \
+      is '\(variableValue)'.
+      """)
+    }
+  }
+}

FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift

@@ -243,25 +243,29 @@ final class IntegrationTests: XCTestCase {
   // MARK: - Imagen
 
   func testGenerateImage_inlineData() async throws {
+    try IntegrationTestUtils.skipUnless(environmentVariable: "VTXIntegrationImagen")
     let imagePrompt = """
     A realistic photo of a male lion, mane thick and dark, standing proudly on a rocky outcrop
     overlooking a vast African savanna at sunset. Golden hour light, long shadows, sharp focus on
     the lion, shallow depth of field, detailed fur texture, DSLR, 85mm lens.
     """
 
-    let imageResponse = try await imagenModel.generateImages(prompt: imagePrompt)
-
-    XCTAssertNil(imageResponse.raiFilteredReason)
-    XCTAssertEqual(imageResponse.images.count, 1)
-    let image = try XCTUnwrap(imageResponse.images.first)
-
-    let textResponse = try await model.generateContent(
-      InlineDataPart(data: image.data, mimeType: "image/png"),
-      "What is the name of this animal? Answer with the animal name only."
-    )
-
-    let text = try XCTUnwrap(textResponse.text).trimmingCharacters(in: .whitespacesAndNewlines)
-    XCTAssertEqual(text, "Lion")
+    let response = try await imagenModel.generateImages(prompt: imagePrompt)
+
+    XCTAssertNil(response.raiFilteredReason)
+    XCTAssertEqual(response.images.count, 1)
+    let image = try XCTUnwrap(response.images.first)
+    XCTAssertEqual(image.mimeType, "image/png")
+    XCTAssertGreaterThan(image.data.count, 0)
+    let imagenImage = image.imagenImage
+    XCTAssertEqual(imagenImage.mimeType, image.mimeType)
+    XCTAssertEqual(imagenImage.bytesBase64Encoded, image.data.base64EncodedString())
+    XCTAssertNil(imagenImage.gcsURI)
+    #if canImport(UIKit)
+      let uiImage = try XCTUnwrap(UIImage(data: image.data))
+      XCTAssertEqual(uiImage.size.width, 1024.0)
+      XCTAssertEqual(uiImage.size.height, 1024.0)
+    #endif
   }
 }
 

FirebaseVertexAI/Tests/TestApp/VertexAITestApp.xcodeproj/project.pbxproj

@@ -7,6 +7,7 @@
 	objects = {
 
 /* Begin PBXBuildFile section */
+		862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */ = {isa = PBXBuildFile; fileRef = 862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */; };
 		8661385C2CC943DD00F4B78E /* TestApp.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8661385B2CC943DD00F4B78E /* TestApp.swift */; };
 		8661385E2CC943DD00F4B78E /* ContentView.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8661385D2CC943DD00F4B78E /* ContentView.swift */; };
 		8661386E2CC943DE00F4B78E /* IntegrationTests.swift in Sources */ = {isa = PBXBuildFile; fileRef = 8661386D2CC943DE00F4B78E /* IntegrationTests.swift */; };
@@ -33,6 +34,7 @@
 /* End PBXContainerItemProxy section */
 
 /* Begin PBXFileReference section */
+		862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = IntegrationTestUtils.swift; sourceTree = "<group>"; };
 		866138582CC943DD00F4B78E /* VertexAITestApp-SPM.app */ = {isa = PBXFileReference; explicitFileType = wrapper.application; includeInIndex = 0; path = "VertexAITestApp-SPM.app"; sourceTree = BUILT_PRODUCTS_DIR; };
 		8661385B2CC943DD00F4B78E /* TestApp.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = TestApp.swift; sourceTree = "<group>"; };
 		8661385D2CC943DD00F4B78E /* ContentView.swift */ = {isa = PBXFileReference; lastKnownFileType = sourcecode.swift; path = ContentView.swift; sourceTree = "<group>"; };
@@ -122,6 +124,7 @@
 			children = (
 				868A7C4D2CCC1F4700E449DD /* Credentials.swift */,
 				8661386D2CC943DE00F4B78E /* IntegrationTests.swift */,
+				862218802D04E08D007ED2D4 /* IntegrationTestUtils.swift */,
 			);
 			path = Integration;
 			sourceTree = "<group>";
@@ -265,6 +268,7 @@
 			files = (
 				8698D7462CD3CF3600ABA833 /* FirebaseAppTestUtils.swift in Sources */,
 				868A7C4F2CCC229F00E449DD /* Credentials.swift in Sources */,
+				862218812D04E098007ED2D4 /* IntegrationTestUtils.swift in Sources */,
 				8661386E2CC943DE00F4B78E /* IntegrationTests.swift in Sources */,
 			);
 			runOnlyForDeploymentPostprocessing = 0;