[Vertex AI] Replace ImagenImage protocol with _ImagenImage struct (…

…#14249)
firebase · Dec 12, 2024 · 1de7f6f · 1de7f6f
1 parent c986629
commit 1de7f6f
Show file tree

Hide file tree

Showing 8 changed files with 15 additions and 55 deletions.
diff --git a/FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift b/FirebaseVertexAI/Sources/Types/Internal/Imagen/InternalImagenImage.swift
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenFileDataImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenFileDataImage.swift
@@ -27,8 +27,8 @@ public struct ImagenFileDataImage {
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension ImagenFileDataImage: ImagenImageRepresentable {
-  public var imagenImage: any ImagenImage {
-    InternalImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)
+  public var _imagenImage: _ImagenImage {
+    _ImagenImage(mimeType: mimeType, bytesBase64Encoded: nil, gcsURI: gcsURI)
   }
 }
 

diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenImage.swift
@@ -15,8 +15,8 @@
 import Foundation
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public protocol ImagenImage: ImagenImageRepresentable {
-  var mimeType: String { get }
-  var bytesBase64Encoded: String? { get }
-  var gcsURI: String? { get }
+public struct _ImagenImage {
+  let mimeType: String
+  let bytesBase64Encoded: String?
+  let gcsURI: String?
 }
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenImageRepresentable.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenImageRepresentable.swift
@@ -16,12 +16,5 @@ import Foundation
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 public protocol ImagenImageRepresentable {
-  var imagenImage: any ImagenImage { get }
-}
-
-@available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
-public extension ImagenImage {
-  var imagenImage: any ImagenImage {
-    return self
-  }
+  var _imagenImage: _ImagenImage { get }
 }
diff --git a/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineDataImage.swift b/FirebaseVertexAI/Sources/Types/Public/Imagen/ImagenInlineDataImage.swift
@@ -31,8 +31,8 @@ public struct ImagenInlineDataImage {
 
 @available(iOS 15.0, macOS 12.0, macCatalyst 15.0, tvOS 15.0, watchOS 8.0, *)
 extension ImagenInlineDataImage: ImagenImageRepresentable {
-  public var imagenImage: any ImagenImage {
-    InternalImagenImage(
+  public var _imagenImage: _ImagenImage {
+    _ImagenImage(
       mimeType: mimeType,
       bytesBase64Encoded: data.base64EncodedString(),
       gcsURI: nil

diff --git a/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift b/FirebaseVertexAI/Tests/TestApp/Tests/Integration/IntegrationTests.swift
@@ -267,10 +267,6 @@ final class IntegrationTests: XCTestCase {
     let image = try XCTUnwrap(response.images.first)
     XCTAssertEqual(image.mimeType, "image/jpeg")
     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, 1408.0)

diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenFileDataImageTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenFileDataImageTests.swift
@@ -35,9 +35,9 @@ final class ImagenFileDataImageTests: XCTestCase {
 
     XCTAssertEqual(image.mimeType, mimeType)
     XCTAssertEqual(image.gcsURI, gcsURI)
-    XCTAssertEqual(image.imagenImage.mimeType, mimeType)
-    XCTAssertEqual(image.imagenImage.gcsURI, gcsURI)
-    XCTAssertNil(image.imagenImage.bytesBase64Encoded)
+    XCTAssertEqual(image._imagenImage.mimeType, mimeType)
+    XCTAssertEqual(image._imagenImage.gcsURI, gcsURI)
+    XCTAssertNil(image._imagenImage.bytesBase64Encoded)
   }
 
   func testDecodeImage_missingGCSURI_throws() throws {

diff --git a/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineDataImageTests.swift b/FirebaseVertexAI/Tests/Unit/Types/Imagen/ImagenInlineDataImageTests.swift
@@ -35,9 +35,9 @@ final class ImagenInlineDataImageTests: XCTestCase {
 
     XCTAssertEqual(image.mimeType, mimeType)
     XCTAssertEqual(image.data.base64EncodedString(), bytesBase64Encoded)
-    XCTAssertEqual(image.imagenImage.mimeType, mimeType)
-    XCTAssertEqual(image.imagenImage.bytesBase64Encoded, bytesBase64Encoded)
-    XCTAssertNil(image.imagenImage.gcsURI)
+    XCTAssertEqual(image._imagenImage.mimeType, mimeType)
+    XCTAssertEqual(image._imagenImage.bytesBase64Encoded, bytesBase64Encoded)
+    XCTAssertNil(image._imagenImage.gcsURI)
   }
 
   func testDecodeImage_missingBytesBase64Encoded_throws() throws {