Fix tests

Change-Id: Ifa610965c5d6c38123080a7e16416ac325418285

Author: MarkDaoust

google/generativeai/types/generation_types.py

@@ -321,6 +321,7 @@ def _join_code_execution_result(result_1, result_2):
 
 
 def _join_candidates(candidates: Iterable[protos.Candidate]):
+    """Joins stream chunks of a single candidate."""
     candidates = tuple(candidates)
 
     index = candidates[0].index  # These should all be the same.
@@ -336,6 +337,7 @@ def _join_candidates(candidates: Iterable[protos.Candidate]):
 
 
 def _join_candidate_lists(candidate_lists: Iterable[list[protos.Candidate]]):
+    """Joins stream chunks where each chunk is a list of candidate chunks."""
     # Assuming that is a candidate ends, it is no longer returned in the list of
     # candidates and that's why candidates have an index
     candidates = collections.defaultdict(list)
@@ -359,10 +361,16 @@ def _join_prompt_feedbacks(
 
 def _join_chunks(chunks: Iterable[protos.GenerateContentResponse]):
     chunks = tuple(chunks)
+    if 'usage_metadata' in chunks[-1]:
+        usage_metadata = chunks[-1].usage_metadata
+    else:
+        usage_metadata=None
+
+
     return protos.GenerateContentResponse(
         candidates=_join_candidate_lists(c.candidates for c in chunks),
         prompt_feedback=_join_prompt_feedbacks(c.prompt_feedback for c in chunks),
-        usage_metadata=chunks[-1].usage_metadata,
+        usage_metadata=usage_metadata,
     )
 
 

tests/test_generation.py

@@ -1,4 +1,20 @@
+# -*- coding: utf-8 -*-
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import inspect
+import json
 import string
 import textwrap
 from typing_extensions import TypedDict
@@ -22,6 +38,7 @@ class Person(TypedDict):
 
 
 class UnitTests(parameterized.TestCase):
+    maxDiff = None
     @parameterized.named_parameters(
         [
             "protos.GenerationConfig",
@@ -416,24 +433,16 @@ def test_join_prompt_feedbacks(self):
                 ],
                 "role": "assistant",
             },
-            "citation_metadata": {"citation_sources": []},
             "index": 0,
-            "finish_reason": 0,
-            "safety_ratings": [],
-            "token_count": 0,
-            "grounding_attributions": [],
+            "citation_metadata": {},
         },
         {
             "content": {
                 "parts": [{"text": "Tell me a story about a magic backpack"}],
                 "role": "assistant",
             },
             "index": 1,
-            "citation_metadata": {"citation_sources": []},
-            "finish_reason": 0,
-            "safety_ratings": [],
-            "token_count": 0,
-            "grounding_attributions": [],
+            "citation_metadata": {},
         },
         {
             "content": {
@@ -458,17 +467,13 @@ def test_join_prompt_feedbacks(self):
                     },
                 ]
             },
-            "finish_reason": 0,
-            "safety_ratings": [],
-            "token_count": 0,
-            "grounding_attributions": [],
         },
     ]
 
     def test_join_candidates(self):
         candidate_lists = [[protos.Candidate(c) for c in cl] for cl in self.CANDIDATE_LISTS]
         result = generation_types._join_candidate_lists(candidate_lists)
-        self.assertEqual(self.MERGED_CANDIDATES, [type(r).to_dict(r) for r in result])
+        self.assertEqual(self.MERGED_CANDIDATES, [type(r).to_dict(r, including_default_value_fields=False) for r in result])
 
     def test_join_chunks(self):
         chunks = [protos.GenerateContentResponse(candidates=cl) for cl in self.CANDIDATE_LISTS]
@@ -480,6 +485,8 @@ def test_join_chunks(self):
             ],
         )
 
+        chunks[-1].usage_metadata = protos.GenerateContentResponse.UsageMetadata(prompt_token_count=5)
+
         result = generation_types._join_chunks(chunks)
 
         expected = protos.GenerateContentResponse(
@@ -495,10 +502,17 @@ def test_join_chunks(self):
                         }
                     ],
                 },
+                "usage_metadata": {
+                         "prompt_token_count": 5
+                }
+
             },
         )
 
-        self.assertEqual(type(expected).to_dict(expected), type(result).to_dict(expected))
+        expected = json.dumps(type(expected).to_dict(expected, including_default_value_fields=False), indent=4)
+        result = json.dumps(type(result).to_dict(result, including_default_value_fields=False), indent=4)
+
+        self.assertEqual(expected, result)
 
     def test_generate_content_response_iterator_end_to_end(self):
         chunks = [protos.GenerateContentResponse(candidates=cl) for cl in self.CANDIDATE_LISTS]

tests/test_generative_models.py

@@ -935,8 +935,7 @@ def test_repr_for_streaming_start_to_finish(self):
                       "citation_metadata": {}
                     }
                   ],
-                  "prompt_feedback": {},
-                  "usage_metadata": {}
+                  "prompt_feedback": {}
                 }),
             )"""
         )
@@ -964,8 +963,7 @@ def test_repr_for_streaming_start_to_finish(self):
                       "citation_metadata": {}
                     }
                   ],
-                  "prompt_feedback": {},
-                  "usage_metadata": {}
+                  "prompt_feedback": {}
                 }),
             )"""
         )
@@ -1056,8 +1054,7 @@ def no_throw():
                       "citation_metadata": {}
                     }
                   ],
-                  "prompt_feedback": {},
-                  "usage_metadata": {}
+                  "prompt_feedback": {}
                 }),
             ),
             error=ValueError()"""