From 1b6853779c9a2d132876e93a19df05b13824ecb8 Mon Sep 17 00:00:00 2001
From: Nate Bosch <nbosch@google.com>
Date: Wed, 22 May 2024 01:39:57 +0000
Subject: [PATCH] Update test expectations

---
 .../test/generative_model_test.dart           | 25 ++++++++++++-------
 1 file changed, 16 insertions(+), 9 deletions(-)

diff --git a/pkgs/google_generative_ai/test/generative_model_test.dart b/pkgs/google_generative_ai/test/generative_model_test.dart
index 1e69053..3a58b3b 100644
--- a/pkgs/google_generative_ai/test/generative_model_test.dart
+++ b/pkgs/google_generative_ai/test/generative_model_test.dart
@@ -99,6 +99,7 @@ void main() {
               ),
             );
             expect(request, {
+              'model': 'models/$defaultModelName',
               'contents': [
                 {
                   'role': 'user',
@@ -311,6 +312,7 @@ void main() {
               ),
             );
             expect(request, {
+              'model': 'models/$defaultModelName',
               'contents': [
                 {
                   'role': 'user',
@@ -415,14 +417,17 @@ void main() {
               ),
             );
             expect(request, {
-              'contents': [
-                {
-                  'role': 'user',
-                  'parts': [
-                    {'text': prompt},
-                  ],
-                },
-              ],
+              'generateContentRequest': {
+                'model': 'models/$defaultModelName',
+                'contents': [
+                  {
+                    'role': 'user',
+                    'parts': [
+                      {'text': prompt},
+                    ],
+                  },
+                ],
+              }
             });
           },
           response: {'totalTokens': 2},
@@ -460,7 +465,9 @@ void main() {
               ),
             ),
           ),
-          verifyRequest: (_, request) {
+          verifyRequest: (_, countTokensRequest) {
+            final request = countTokensRequest['generateContentRequest']
+                as Map<String, Object?>;
             expect(request['safetySettings'], [
               {
                 'category': 'HARM_CATEGORY_DANGEROUS_CONTENT',