Skip to content

Commit

Permalink
Add ModelContent tests for fileData parts (google-gemini#163)
Browse files Browse the repository at this point in the history
  • Loading branch information
andrewheard authored and G.Dev.Ssomsak committed Jun 21, 2024
1 parent d6b4250 commit 25e222a
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 0 deletions.
4 changes: 4 additions & 0 deletions Sources/GoogleAI/ModelContent.swift
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,10 @@ public struct ModelContent: Equatable {

/// URI-based data with a specified media type.
///
/// > Important: Files must be uploaded using the
/// > [`media.upload` REST API](https://ai.google.dev/api/rest/v1beta/media/upload) or another
/// > Gemini SDK.
///
/// > Note: Supported media types depends on the model; see
/// > [supported file
/// > formats](https://ai.google.dev/tutorials/prompting_with_media#supported_file_formats)
Expand Down
7 changes: 7 additions & 0 deletions Tests/GoogleAITests/GoogleAITests.swift
Original file line number Diff line number Diff line change
Expand Up @@ -91,6 +91,13 @@ final class GoogleGenerativeAITests: XCTestCase {
let _ = try await genAI.generateContent(str)
let _ = try await genAI.generateContent([str])
let _ = try await genAI.generateContent(str, "abc", "def")
let _ = try await genAI.generateContent(
str,
ModelContent.Part.fileData(
mimetype: "image/jpeg",
uri: "https://generativelanguage.googleapis.com/v1beta/files/rand0mha5sh"
)
)
#if canImport(UIKit)
_ = try await genAI.generateContent(UIImage())
_ = try await genAI.generateContent([UIImage()])
Expand Down
48 changes: 48 additions & 0 deletions Tests/GoogleAITests/ModelContentTests.swift
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
// 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 GoogleGenerativeAI
import XCTest

@available(iOS 15.0, macOS 11.0, macCatalyst 15.0, *)
final class ModelContentTests: XCTestCase {
let encoder = JSONEncoder()

override func setUp() {
encoder.outputFormatting = .init(
arrayLiteral: .prettyPrinted, .sortedKeys, .withoutEscapingSlashes
)
}

// MARK: ModelContent.Part Encoding

func testEncodeFileDataPart() throws {
let mimeType = "image/jpeg"
let fileURI = "gs://test-bucket/image.jpg"
let fileDataPart = ModelContent.Part.fileData(mimetype: mimeType, uri: fileURI)

let jsonData = try encoder.encode(fileDataPart)

let json = try XCTUnwrap(String(data: jsonData, encoding: .utf8))
XCTAssertEqual(json, """
{
"fileData" : {
"file_uri" : "\(fileURI)",
"mime_type" : "\(mimeType)"
}
}
""")
}
}

0 comments on commit 25e222a

Please # to comment.