Skip to content

Commit 3f49b09

Browse files
paulb777paulb777
committed
streamline integration tests
1 parent a49dee1 commit 3f49b09

File tree

2 files changed

+74
-104
lines changed

2 files changed

+74
-104
lines changed

‎FirebaseAI/Sources/TemplateGenerateContentRequest.swift‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@ struct TemplateGenerateContentRequest: Sendable {
2121
let history: [ModelContent]
2222
let projectID: String
2323
let stream: Bool
24-
2524
let apiConfig: APIConfig
2625
let options: RequestOptions
2726
}
@@ -53,7 +52,6 @@ extension TemplateGenerateContentRequest: GenerativeAIRequest {
5352
let templateName = template.hasSuffix(".prompt") ? template : "\(template).prompt"
5453
urlString += "/templates/\(templateName):templateGenerateContent"
5554
if stream {
56-
// TODO: Fix this.
5755
urlString += "?alt=sse"
5856
}
5957
return URL(string: urlString)!

‎FirebaseAI/Tests/TestApp/Tests/Integration/ServerPromptTemplateIntegrationTests.swift‎

Lines changed: 74 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -25,59 +25,47 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
2525
func testGenerateContentWithText() async throws {
2626
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
2727
let userName = "paul"
28-
do {
29-
let response = try await model.generateContent(
30-
template: "greeting",
31-
variables: [
32-
"name": userName,
33-
"language": "Spanish",
34-
]
35-
)
36-
let text = try XCTUnwrap(response.text)
37-
print(text)
38-
XCTAssert(text.contains("Paul"))
39-
} catch {
40-
XCTFail("An error occurred: \(error)")
41-
}
28+
let response = try await model.generateContent(
29+
template: "greeting",
30+
variables: [
31+
"name": userName,
32+
"language": "Spanish",
33+
]
34+
)
35+
let text = try XCTUnwrap(response.text)
36+
print(text)
37+
XCTAssert(text.contains("Paul"))
4238
}
4339

4440
func testGenerateContentStream() async throws {
4541
let model = FirebaseAI.firebaseAI(backend: .vertexAI()).templateGenerativeModel()
4642
let userName = "paul"
47-
do {
48-
let stream = try model.generateContentStream(
49-
template: "greeting",
50-
variables: [
51-
"name": userName,
52-
"language": "English",
53-
]
54-
)
55-
var resultText = ""
56-
for try await response in stream {
57-
if let text = response.text {
58-
resultText += text
59-
}
43+
let stream = try model.generateContentStream(
44+
template: "greeting",
45+
variables: [
46+
"name": userName,
47+
"language": "English",
48+
]
49+
)
50+
var resultText = ""
51+
for try await response in stream {
52+
if let text = response.text {
53+
resultText += text
6054
}
61-
XCTAssert(resultText.contains("Paul"))
62-
} catch {
63-
XCTFail("An error occurred: \(error)")
6455
}
56+
XCTAssert(resultText.contains("Paul"))
6557
}
6658

6759
func testGenerateImages() async throws {
6860
let imagenModel = FirebaseAI.firebaseAI(backend: .vertexAI()).templateImagenModel()
6961
let imagenPrompt = "A cat picture"
70-
do {
71-
let response = try await imagenModel.generateImages(
72-
template: "generate_images",
73-
variables: [
74-
"prompt": imagenPrompt,
75-
]
76-
)
77-
XCTAssertEqual(response.images.count, 3)
78-
} catch {
79-
XCTFail("An error occurred: \(error)")
80-
}
62+
let response = try await imagenModel.generateImages(
63+
template: "generate_images",
64+
variables: [
65+
"prompt": imagenPrompt,
66+
]
67+
)
68+
XCTAssertEqual(response.images.count, 3)
8169
}
8270

8371
func testGenerateContentWithMedia() async throws {
@@ -86,21 +74,17 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
8674
if let imageBytes = image.jpegData(compressionQuality: 0.8) {
8775
let base64Image = imageBytes.base64EncodedString()
8876

89-
do {
90-
let response = try await model.generateContent(
91-
template: "media",
92-
variables: [
93-
"imageData": [
94-
"isInline": true,
95-
"mimeType": "image/jpeg",
96-
"contents": base64Image,
97-
],
98-
]
99-
)
100-
XCTAssert(response.text?.isEmpty == false)
101-
} catch {
102-
XCTFail("An error occurred: \(error)")
103-
}
77+
let response = try await model.generateContent(
78+
template: "media",
79+
variables: [
80+
"imageData": [
81+
"isInline": true,
82+
"mimeType": "image/jpeg",
83+
"contents": base64Image,
84+
],
85+
]
86+
)
87+
XCTAssert(response.text?.isEmpty == false)
10488
} else {
10589
XCTFail("Could not get image data.")
10690
}
@@ -112,27 +96,23 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
11296
if let imageBytes = image.jpegData(compressionQuality: 0.8) {
11397
let base64Image = imageBytes.base64EncodedString()
11498

115-
do {
116-
let stream = try model.generateContentStream(
117-
template: "media",
118-
variables: [
119-
"imageData": [
120-
"isInline": true,
121-
"mimeType": "image/jpeg",
122-
"contents": base64Image,
123-
],
124-
]
125-
)
126-
var resultText = ""
127-
for try await response in stream {
128-
if let text = response.text {
129-
resultText += text
130-
}
99+
let stream = try model.generateContentStream(
100+
template: "media",
101+
variables: [
102+
"imageData": [
103+
"isInline": true,
104+
"mimeType": "image/jpeg",
105+
"contents": base64Image,
106+
],
107+
]
108+
)
109+
var resultText = ""
110+
for try await response in stream {
111+
if let text = response.text {
112+
resultText += text
131113
}
132-
XCTAssert(resultText.isEmpty == false)
133-
} catch {
134-
XCTFail("An error occurred: \(error)")
135114
}
115+
XCTAssert(resultText.isEmpty == false)
136116
} else {
137117
XCTFail("Could not get image data.")
138118
}
@@ -148,17 +128,13 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
148128

149129
let userMessage = "What's the weather like?"
150130

151-
do {
152-
let response = try await chatSession.sendMessage(
153-
userMessage,
154-
variables: ["message": userMessage]
155-
)
156-
XCTAssert(response.text?.isEmpty == false)
157-
XCTAssertEqual(chatSession.history.count, 4)
158-
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
159-
} catch {
160-
XCTFail("An error occurred: \(error)")
161-
}
131+
let response = try await chatSession.sendMessage(
132+
userMessage,
133+
variables: ["message": userMessage]
134+
)
135+
XCTAssert(response.text?.isEmpty == false)
136+
XCTAssertEqual(chatSession.history.count, 4)
137+
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
162138
}
163139

164140
func testChatStream() async throws {
@@ -171,22 +147,18 @@ final class ServerPromptTemplateIntegrationTests: XCTestCase {
171147

172148
let userMessage = "What's the weather like?"
173149

174-
do {
175-
let stream = try chatSession.sendMessageStream(
176-
userMessage,
177-
variables: ["message": userMessage]
178-
)
179-
var resultText = ""
180-
for try await response in stream {
181-
if let text = response.text {
182-
resultText += text
183-
}
150+
let stream = try chatSession.sendMessageStream(
151+
userMessage,
152+
variables: ["message": userMessage]
153+
)
154+
var resultText = ""
155+
for try await response in stream {
156+
if let text = response.text {
157+
resultText += text
184158
}
185-
XCTAssert(resultText.isEmpty == false)
186-
XCTAssertEqual(chatSession.history.count, 4)
187-
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
188-
} catch {
189-
XCTFail("An error occurred: \(error)")
190159
}
160+
XCTAssert(resultText.isEmpty == false)
161+
XCTAssertEqual(chatSession.history.count, 4)
162+
XCTAssertEqual((chatSession.history[2].parts.first as? TextPart)?.text, userMessage)
191163
}
192-
}
164+
}

0 commit comments

Comments
 (0)