From c975cb0dbd3cc01a0edba00369f452aa5876c9a7 Mon Sep 17 00:00:00 2001 From: NatashaTheRobot Date: Sat, 21 Dec 2024 21:01:23 +0530 Subject: [PATCH] added files tests --- ...emini.APISpecification.RequestBodies.swift | 7 ++ .../API/_Gemini.APISpecification.swift | 17 +++ .../Intramodular/Models/_Gemini.File.swift | 7 ++ .../Intramodular/_Gemini.Client+Files.swift | 16 +++ .../Intramodular/_GeminiTests+Files.swift | 100 ++++++++++++++++++ 5 files changed, 147 insertions(+) create mode 100644 Tests/_Gemini/Intramodular/_GeminiTests+Files.swift diff --git a/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.RequestBodies.swift b/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.RequestBodies.swift index 8d84078..7003aeb 100644 --- a/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.RequestBodies.swift +++ b/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.RequestBodies.swift @@ -215,6 +215,13 @@ extension _Gemini.APISpecification { public let name: _Gemini.File.Name } + public struct FileListInput: Codable { + // Maximum number of Files to return per page. If unspecified, defaults to 10. Maximum pageSize is 100. + public let pageSize: Int? + // A page token from a previous files.list call. + public let pageToken: String? + } + // Fine Tuning public struct CreateTunedModel: Codable { diff --git a/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.swift b/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.swift index e893ecf..fe289d0 100644 --- a/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.swift +++ b/Sources/_Gemini/Intramodular/API/_Gemini.APISpecification.swift @@ -88,6 +88,23 @@ extension _Gemini { }) var getFile = Endpoint() + @GET + @Path("/v1beta/files") + @Query({ context -> [String : String] in + var parameters: [String : String] = [:] + + if let pageSize = context.input.pageSize { + parameters["pageSize"] = String(pageSize) + } + + if let pageToken = context.input.pageToken { + parameters["pageToken"] = pageToken + } + + return parameters + }) + var listFiles = Endpoint() + // Delete File endpoint @DELETE @Path({ context -> String in diff --git a/Sources/_Gemini/Intramodular/Models/_Gemini.File.swift b/Sources/_Gemini/Intramodular/Models/_Gemini.File.swift index 4fd7c45..077ed26 100644 --- a/Sources/_Gemini/Intramodular/Models/_Gemini.File.swift +++ b/Sources/_Gemini/Intramodular/Models/_Gemini.File.swift @@ -8,6 +8,7 @@ import Foundation extension _Gemini { + public struct File: Codable { public let createTime: String? public let expirationTime: String? @@ -29,6 +30,12 @@ extension _Gemini { public let videoDuration: String } } + + public struct FileList: Codable { + public let files: [_Gemini.File] + // A token that can be sent as a pageToken into a subsequent files.list call. + public let nextPageToken: String + } } extension _Gemini.File { diff --git a/Sources/_Gemini/Intramodular/_Gemini.Client+Files.swift b/Sources/_Gemini/Intramodular/_Gemini.Client+Files.swift index ff454ba..e1d5d38 100644 --- a/Sources/_Gemini/Intramodular/_Gemini.Client+Files.swift +++ b/Sources/_Gemini/Intramodular/_Gemini.Client+Files.swift @@ -9,6 +9,7 @@ import NetworkKit import Swallow extension _Gemini.Client { + public func uploadFile( fileData: Data, mimeType: HTTPMediaType, @@ -48,6 +49,21 @@ extension _Gemini.Client { } } + public func listFiles( + pageSize: Int? = nil, + pageToken: String? = nil + ) async throws -> _Gemini.FileList { + do { + let input = _Gemini.APISpecification.RequestBodies.FileListInput( + pageSize: pageSize, + pageToken: pageToken + ) + return try await run(\.listFiles, with: input) + } catch { + throw _Gemini.APIError.unknown(message: "Failed to get file status: \(error.localizedDescription)") + } + } + public func deleteFile( fileURL: URL ) async throws { diff --git a/Tests/_Gemini/Intramodular/_GeminiTests+Files.swift b/Tests/_Gemini/Intramodular/_GeminiTests+Files.swift new file mode 100644 index 0000000..85a6d7e --- /dev/null +++ b/Tests/_Gemini/Intramodular/_GeminiTests+Files.swift @@ -0,0 +1,100 @@ +// +// _GeminiTests+Files.swift +// AI +// +// Created by Natasha Murashev on 12/21/24. +// + +import Testing +import Foundation +import _Gemini + +@Suite struct GeminiFileTests { + + private let fileURL = URL(string: "https://upload.wikimedia.org/wikipedia/en/7/77/EricCartman.png")! + private var fileData: Data? = nil + private var fileName: _Gemini.File.Name? = nil + + @Test mutating func testUploadFile() async throws { + let file: _Gemini.File = try await uploadFile() + + print(file) + #expect(file.name.isNotNil) + #expect(((file.name?.rawValue.starts(with: "files/")) == true)) + } + + @Test mutating func testGetFile() async throws { + + if fileName == nil { + fileName = try await uploadFile().name + } + + guard let fileName = fileName else { + #expect(Bool(false), "The file name is invalid") + return + } + + let file: _Gemini.File = try await client.getFile(name: fileName) + #expect(file.name?.rawValue == fileName.rawValue) + } + + @Test mutating func testListFiles() async throws { + let file = try await uploadFile() + let fileList: _Gemini.FileList = try await client.listFiles() + let files: [_Gemini.File] = fileList.files + + guard let fileName = file.name else { + #expect(Bool(false), "The uploaded file has no valid name.") + return + } + + let uploadedFileIsPresent = files.contains { $0.name! == fileName } + #expect(uploadedFileIsPresent, "Expected the newly uploaded file to be in the returned file list.") + } + + @Test mutating func testDeleteFile() async throws { + let uploadedFile = try await uploadFile() + let file: _Gemini.File = try await client.getFile(name: uploadedFile.name!) + + try await client.deleteFile(fileURL: file.uri) + do { + let _ = try await client.getFile(name: uploadedFile.name!) + #expect(Bool(false), "Expected getFile to throw when fetching a deleted file, but it succeeded.") + } catch { + #expect(Bool(true), "getFile threw an error, as expected, when trying to retrieve a deleted file.") + } + } +} + +extension GeminiFileTests { + + private mutating func uploadFile() async throws -> _Gemini.File { + let data = try await downloadFile(from: fileURL) + + let file = try await client.uploadFile( + fileData: data, + mimeType: .custom("image/png"), + displayName: UUID().uuidString + ) + + self.fileName = file.name + + return file + } + + private func downloadFile(from url: URL) async throws -> Data { + if let data = fileData { + return data + } + + let (data, response) = try await URLSession.shared.data(from: url) + + guard let httpResponse = response as? HTTPURLResponse, + (200...299).contains(httpResponse.statusCode) else { + throw _Gemini.APIError.unknown(message: "Failed to download file from URL") + } + + return data + } +} +