File tree 2 files changed +17
-1
lines changed
src/test/java/com/google/firebase/ai
2 files changed +17
-1
lines changed Original file line number Diff line number Diff line change @@ -24,6 +24,7 @@ import com.google.firebase.ai.util.goldenDevAPIUnaryFile
24
24
import io.kotest.assertions.throwables.shouldThrow
25
25
import io.kotest.matchers.should
26
26
import io.kotest.matchers.shouldBe
27
+ import io.kotest.matchers.shouldNotBe
27
28
import io.ktor.http.HttpStatusCode
28
29
import kotlin.time.Duration.Companion.seconds
29
30
import kotlinx.coroutines.withTimeout
@@ -56,6 +57,21 @@ internal class DevAPIUnarySnapshotTests {
56
57
}
57
58
}
58
59
60
+ @Test
61
+ fun `citation returns correctly` () =
62
+ goldenDevAPIUnaryFile(" unary-success-citations.json" ) {
63
+ withTimeout(testTimeout) {
64
+ val response = model.generateContent(" prompt" )
65
+
66
+ response.candidates.isEmpty() shouldBe false
67
+ response.candidates.first().citationMetadata?.citations?.size shouldBe 4
68
+ response.candidates.first().citationMetadata?.citations?.forEach {
69
+ it.startIndex shouldNotBe null
70
+ it.endIndex shouldNotBe null
71
+ }
72
+ }
73
+ }
74
+
59
75
@Test
60
76
fun `response blocked for safety` () =
61
77
goldenDevAPIUnaryFile(" unary-failure-finish-reason-safety.txt" ) {
Original file line number Diff line number Diff line change 17
17
# This script replaces mock response files for Vertex AI unit tests with a fresh
18
18
# clone of the shared repository of Vertex AI test data.
19
19
20
- RESPONSES_VERSION=' v10 .*' # The major version of mock responses to use
20
+ RESPONSES_VERSION=' v11 .*' # The major version of mock responses to use
21
21
REPO_NAME=" vertexai-sdk-test-data"
22
22
REPO_LINK=" https://github.com/FirebaseExtended/$REPO_NAME .git"
23
23
You can’t perform that action at this time.
0 commit comments