From 7a9ceeebecb0a081e98cbd1cb4f4f56699d087f2 Mon Sep 17 00:00:00 2001
From: Hadley Wickham <h.wickham@gmail.com>
Date: Wed, 18 Dec 2024 14:32:36 -0600
Subject: [PATCH] Add a test

---
 tests/testthat/test-provider-gemini.R | 23 +++++++++++++++++++++++
 1 file changed, 23 insertions(+)

diff --git a/tests/testthat/test-provider-gemini.R b/tests/testthat/test-provider-gemini.R
index a2546de0..90d675db 100644
--- a/tests/testthat/test-provider-gemini.R
+++ b/tests/testthat/test-provider-gemini.R
@@ -48,3 +48,26 @@ test_that("can use images", {
   test_images_inline(chat_fun)
   test_images_remote_error(chat_fun)
 })
+
+# chunk merging ----------------------------------------------------------
+
+test_that("can merge text output", {
+  # output from "tell me a joke" with text changed
+  messages <- c(
+    '{"candidates": [{"content": {"parts": [{"text": "a"}],"role": "model"}}],"usageMetadata": {"promptTokenCount": 5,"totalTokenCount": 5},"modelVersion": "gemini-1.5-flash"}',
+    '{"candidates": [{"content": {"parts": [{"text": "b"}],"role": "model"}}],"usageMetadata": {"promptTokenCount": 5,"totalTokenCount": 5},"modelVersion": "gemini-1.5-flash"}',
+    '{"candidates": [{"content": {"parts": [{"text": "c"}],"role": "model"},"finishReason": "STOP"}],"usageMetadata": {"promptTokenCount": 5,"candidatesTokenCount": 17,"totalTokenCount": 22},"modelVersion": "gemini-1.5-flash"}'
+  )
+  chunks <- lapply(messages, jsonlite::parse_json)
+
+  out <- merge_gemini_chunks(chunks[[1]], chunks[[2]])
+  out <- merge_gemini_chunks(out, chunks[[3]])
+
+  expect_equal(out$candidates[[1]]$content$parts[[1]]$text, "abc")
+  expect_equal(out$usageMetadata, list(
+    promptTokenCount = 5,
+    candidatesTokenCount = 17,
+    totalTokenCount = 22
+  ))
+  expect_equal(out$candidates[[1]]$finishReason, "STOP")
+})