@@ -25,25 +25,24 @@ import com.google.firebase.vertexai.type.ResponseStoppedException
25
25
import com.google.firebase.vertexai.type.SerializationException
26
26
import com.google.firebase.vertexai.type.ServerException
27
27
import com.google.firebase.vertexai.type.TextPart
28
- import com.google.firebase.vertexai.util.goldenStreamingFile
28
+ import com.google.firebase.vertexai.util.goldenVertexStreamingFile
29
29
import io.kotest.assertions.throwables.shouldThrow
30
30
import io.kotest.matchers.nulls.shouldNotBeNull
31
31
import io.kotest.matchers.shouldBe
32
32
import io.kotest.matchers.string.shouldContain
33
33
import io.ktor.http.HttpStatusCode
34
34
import kotlin.time.Duration.Companion.seconds
35
35
import kotlinx.coroutines.flow.collect
36
- import kotlinx.coroutines.flow.first
37
36
import kotlinx.coroutines.flow.toList
38
37
import kotlinx.coroutines.withTimeout
39
38
import org.junit.Test
40
39
41
- internal class StreamingSnapshotTests {
40
+ internal class VertexAIStreamingSnapshotTests {
42
41
private val testTimeout = 5 .seconds
43
42
44
43
@Test
45
44
fun `short reply` () =
46
- goldenStreamingFile (" streaming-success-basic-reply-short.txt" ) {
45
+ goldenVertexStreamingFile (" streaming-success-basic-reply-short.txt" ) {
47
46
val responses = model.generateContentStream(" prompt" )
48
47
49
48
withTimeout(testTimeout) {
@@ -57,7 +56,7 @@ internal class StreamingSnapshotTests {
57
56
58
57
@Test
59
58
fun `long reply` () =
60
- goldenStreamingFile (" streaming-success-basic-reply-long.txt" ) {
59
+ goldenVertexStreamingFile (" streaming-success-basic-reply-long.txt" ) {
61
60
val responses = model.generateContentStream(" prompt" )
62
61
63
62
withTimeout(testTimeout) {
@@ -73,7 +72,7 @@ internal class StreamingSnapshotTests {
73
72
74
73
@Test
75
74
fun `unknown enum in safety ratings` () =
76
- goldenStreamingFile (" streaming-success-unknown-safety-enum.txt" ) {
75
+ goldenVertexStreamingFile (" streaming-success-unknown-safety-enum.txt" ) {
77
76
val responses = model.generateContentStream(" prompt" )
78
77
79
78
withTimeout(testTimeout) {
@@ -88,7 +87,7 @@ internal class StreamingSnapshotTests {
88
87
89
88
@Test
90
89
fun `unknown enum in finish reason` () =
91
- goldenStreamingFile (" streaming-failure-unknown-finish-enum.txt" ) {
90
+ goldenVertexStreamingFile (" streaming-failure-unknown-finish-enum.txt" ) {
92
91
val responses = model.generateContentStream(" prompt" )
93
92
94
93
withTimeout(testTimeout) {
@@ -99,7 +98,7 @@ internal class StreamingSnapshotTests {
99
98
100
99
@Test
101
100
fun `quotes escaped` () =
102
- goldenStreamingFile (" streaming-success-quotes-escaped.txt" ) {
101
+ goldenVertexStreamingFile (" streaming-success-quotes-escaped.txt" ) {
103
102
val responses = model.generateContentStream(" prompt" )
104
103
105
104
withTimeout(testTimeout) {
@@ -114,7 +113,7 @@ internal class StreamingSnapshotTests {
114
113
115
114
@Test
116
115
fun `prompt blocked for safety` () =
117
- goldenStreamingFile (" streaming-failure-prompt-blocked-safety.txt" ) {
116
+ goldenVertexStreamingFile (" streaming-failure-prompt-blocked-safety.txt" ) {
118
117
val responses = model.generateContentStream(" prompt" )
119
118
120
119
withTimeout(testTimeout) {
@@ -125,7 +124,7 @@ internal class StreamingSnapshotTests {
125
124
126
125
@Test
127
126
fun `prompt blocked for safety with message` () =
128
- goldenStreamingFile (" streaming-failure-prompt-blocked-safety-with-message.txt" ) {
127
+ goldenVertexStreamingFile (" streaming-failure-prompt-blocked-safety-with-message.txt" ) {
129
128
val responses = model.generateContentStream(" prompt" )
130
129
131
130
withTimeout(testTimeout) {
@@ -137,23 +136,26 @@ internal class StreamingSnapshotTests {
137
136
138
137
@Test
139
138
fun `empty content` () =
140
- goldenStreamingFile (" streaming-failure-empty-content.txt" ) {
139
+ goldenVertexStreamingFile (" streaming-failure-empty-content.txt" ) {
141
140
val responses = model.generateContentStream(" prompt" )
142
141
143
142
withTimeout(testTimeout) { shouldThrow<SerializationException > { responses.collect() } }
144
143
}
145
144
146
145
@Test
147
146
fun `http errors` () =
148
- goldenStreamingFile(" streaming-failure-http-error.txt" , HttpStatusCode .PreconditionFailed ) {
147
+ goldenVertexStreamingFile(
148
+ " streaming-failure-http-error.txt" ,
149
+ HttpStatusCode .PreconditionFailed
150
+ ) {
149
151
val responses = model.generateContentStream(" prompt" )
150
152
151
153
withTimeout(testTimeout) { shouldThrow<ServerException > { responses.collect() } }
152
154
}
153
155
154
156
@Test
155
157
fun `stopped for safety` () =
156
- goldenStreamingFile (" streaming-failure-finish-reason-safety.txt" ) {
158
+ goldenVertexStreamingFile (" streaming-failure-finish-reason-safety.txt" ) {
157
159
val responses = model.generateContentStream(" prompt" )
158
160
159
161
withTimeout(testTimeout) {
@@ -164,7 +166,7 @@ internal class StreamingSnapshotTests {
164
166
165
167
@Test
166
168
fun `citation parsed correctly` () =
167
- goldenStreamingFile (" streaming-success-citations.txt" ) {
169
+ goldenVertexStreamingFile (" streaming-success-citations.txt" ) {
168
170
val responses = model.generateContentStream(" prompt" )
169
171
170
172
withTimeout(testTimeout) {
@@ -177,7 +179,7 @@ internal class StreamingSnapshotTests {
177
179
178
180
@Test
179
181
fun `stopped for recitation` () =
180
- goldenStreamingFile (" streaming-failure-recitation-no-content.txt" ) {
182
+ goldenVertexStreamingFile (" streaming-failure-recitation-no-content.txt" ) {
181
183
val responses = model.generateContentStream(" prompt" )
182
184
183
185
withTimeout(testTimeout) {
@@ -188,39 +190,39 @@ internal class StreamingSnapshotTests {
188
190
189
191
@Test
190
192
fun `image rejected` () =
191
- goldenStreamingFile (" streaming-failure-image-rejected.txt" , HttpStatusCode .BadRequest ) {
193
+ goldenVertexStreamingFile (" streaming-failure-image-rejected.txt" , HttpStatusCode .BadRequest ) {
192
194
val responses = model.generateContentStream(" prompt" )
193
195
194
196
withTimeout(testTimeout) { shouldThrow<ServerException > { responses.collect() } }
195
197
}
196
198
197
199
@Test
198
200
fun `unknown model` () =
199
- goldenStreamingFile (" streaming-failure-unknown-model.txt" , HttpStatusCode .NotFound ) {
201
+ goldenVertexStreamingFile (" streaming-failure-unknown-model.txt" , HttpStatusCode .NotFound ) {
200
202
val responses = model.generateContentStream(" prompt" )
201
203
202
204
withTimeout(testTimeout) { shouldThrow<ServerException > { responses.collect() } }
203
205
}
204
206
205
207
@Test
206
208
fun `invalid api key` () =
207
- goldenStreamingFile (" streaming-failure-api-key.txt" , HttpStatusCode .BadRequest ) {
209
+ goldenVertexStreamingFile (" streaming-failure-api-key.txt" , HttpStatusCode .BadRequest ) {
208
210
val responses = model.generateContentStream(" prompt" )
209
211
210
212
withTimeout(testTimeout) { shouldThrow<InvalidAPIKeyException > { responses.collect() } }
211
213
}
212
214
213
215
@Test
214
216
fun `invalid json` () =
215
- goldenStreamingFile (" streaming-failure-invalid-json.txt" ) {
217
+ goldenVertexStreamingFile (" streaming-failure-invalid-json.txt" ) {
216
218
val responses = model.generateContentStream(" prompt" )
217
219
218
220
withTimeout(testTimeout) { shouldThrow<SerializationException > { responses.collect() } }
219
221
}
220
222
221
223
@Test
222
224
fun `malformed content` () =
223
- goldenStreamingFile (" streaming-failure-malformed-content.txt" ) {
225
+ goldenVertexStreamingFile (" streaming-failure-malformed-content.txt" ) {
224
226
val responses = model.generateContentStream(" prompt" )
225
227
226
228
withTimeout(testTimeout) { shouldThrow<SerializationException > { responses.collect() } }
0 commit comments