Update test expectations

google-gemini · May 22, 2024 · 1b68537 · 1b68537
1 parent 9a79968
commit 1b68537
Showing 1 changed file with 16 additions and 9 deletions.
diff --git 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',