Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions FirebaseAI/Tests/TestApp/Resources/TestApp.entitlements
Original file line number Diff line number Diff line change
Expand Up @@ -6,5 +6,7 @@
<true/>
<key>com.apple.security.network.client</key>
<true/>
<key>keychain-access-groups</key>
<array/>
</dict>
</plist>
Original file line number Diff line number Diff line change
Expand Up @@ -12,19 +12,25 @@
// See the License for the specific language governing permissions and
// limitations under the License.

import XCTest

import FirebaseAI

final class ServerPromptTemplateIntegrationTests: XCTestCase {
override func setUp() {
super.setUp()
continueAfterFailure = false
}

func testGenerateContentWithText() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI(location: "global"))
.templateGenerativeModel()
import Testing
#if canImport(UIKit)
import UIKit
#endif

struct ServerPromptTemplateIntegrationTests {
private static let testConfigs: [InstanceConfig] = [
.vertexAI_v1beta,
.vertexAI_v1beta_global,
]
private static let imageGenerationTestConfigs: [InstanceConfig] = [.vertexAI_v1beta]

@Test(arguments: [
// The "greeting2" template is only available in the `global` location.
InstanceConfig.vertexAI_v1beta_global,
])
func generateContentWithText(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let userName = "paul"
let response = try await model.generateContent(
template: "greeting2",
Expand All @@ -33,12 +39,13 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
"language": "Spanish",
]
)
let text = try XCTUnwrap(response.text)
XCTAssert(text.contains("Paul"))
let text = try #require(response.text)
#expect(text.contains("Paul"))
}

func testGenerateContentStream() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
@Test(arguments: testConfigs)
func generateContentStream(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let userName = "paul"
let stream = try model.generateContentStream(
template: "greeting.prompt",
Expand All @@ -53,25 +60,30 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
resultText += text
}
}
XCTAssert(resultText.contains("Paul"))
#expect(resultText.contains("Paul"))
}

func testGenerateImages() async throws {
let imagenModel = FirebaseAI.firebaseAI(backend: .vertexAI()).templateImagenModel()
@Test(arguments: imageGenerationTestConfigs)
func generateImages(_ config: InstanceConfig) async throws {
let imagenModel = FirebaseAI.componentInstance(config).templateImagenModel()
let imagenPrompt = "A cat picture"
let response = try await imagenModel.generateImages(
template: "generate_images.prompt",
variables: [
"prompt": imagenPrompt,
]
)
XCTAssertEqual(response.images.count, 3)
#expect(response.images.count == 3)
}

func testGenerateContentWithMedia() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
let image = UIImage(systemName: "photo")!
if let imageBytes = image.jpegData(compressionQuality: 0.8) {
#if canImport(UIKit)
@Test(arguments: testConfigs)
func generateContentWithMedia(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let image = UIImage(systemName: "photo")!
let imageBytes = try #require(
image.jpegData(compressionQuality: 0.8), "Could not get image data."
)
let base64Image = imageBytes.base64EncodedString()

let response = try await model.generateContent(
Expand All @@ -84,16 +96,19 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
],
]
)
XCTAssert(response.text?.isEmpty == false)
} else {
XCTFail("Could not get image data.")
let text = try #require(response.text)
#expect(!text.isEmpty)
}
}

func testGenerateContentStreamWithMedia() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
let image = UIImage(systemName: "photo")!
if let imageBytes = image.jpegData(compressionQuality: 0.8) {
#endif // canImport(UIKit)

#if canImport(UIKit)
@Test(arguments: testConfigs)
func generateContentStreamWithMedia(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let image = UIImage(systemName: "photo")!
let imageBytes = try #require(
image.jpegData(compressionQuality: 0.8), "Could not get image data."
)
let base64Image = imageBytes.base64EncodedString()

let stream = try model.generateContentStream(
Expand All @@ -112,14 +127,13 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
resultText += text
}
}
XCTAssert(resultText.isEmpty == false)
} else {
XCTFail("Could not get image data.")
#expect(!resultText.isEmpty)
}
}
#endif // canImport(UIKit)

func testChat() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
@Test(arguments: testConfigs)
func chat(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let initialHistory = [
ModelContent(role: "user", parts: "Hello!"),
ModelContent(role: "model", parts: "Hi there! How can I help?"),
Expand All @@ -132,13 +146,16 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
userMessage,
variables: ["message": userMessage]
)
XCTAssert(response.text?.isEmpty == false)
XCTAssertEqual(chatSession.history.count, 4)
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
let text = try #require(response.text)
#expect(!text.isEmpty)
#expect(chatSession.history.count == 4)
let textPart = try #require(chatSession.history[2].parts.first as? TextPart)
#expect(textPart.text == userMessage)
}

func testChatStream() async throws {
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
@Test(arguments: testConfigs)
func chatStream(_ config: InstanceConfig) async throws {
let model = FirebaseAI.componentInstance(config).templateGenerativeModel()
let initialHistory = [
ModelContent(role: "user", parts: "Hello!"),
ModelContent(role: "model", parts: "Hi there! How can I help?"),
Expand All @@ -157,8 +174,9 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
resultText += text
}
}
XCTAssert(resultText.isEmpty == false)
XCTAssertEqual(chatSession.history.count, 4)
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
#expect(!resultText.isEmpty)
#expect(chatSession.history.count == 4)
let textPart = try #require(chatSession.history[2].parts.first as? TextPart)
#expect(textPart.text == userMessage)
}
}
Loading