chore:update tests

PiperOrigin-RevId: 912791706

Author: MarkDaoust

google/genai/tests/shared/batches/test_create_delete.py

@@ -17,6 +17,7 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 def create_delete(client, parameters):
@@ -35,7 +36,7 @@ def create_delete(client, parameters):
     pytest_helper.TestTableItem(
         name="test_create_delete_mldev",
         parameters=genai_types._CreateBatchJobParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_FLASH,
             src=[{
                 "contents": [{
                     "parts": [{"text": "Why is the sky blue?"}],

google/genai/tests/shared/batches/test_create_get_cancel.py

@@ -18,6 +18,7 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 def create_get_cancel(client, parameters):
@@ -34,7 +35,7 @@ def create_get_cancel(client, parameters):
     pytest_helper.TestTableItem(
         name="test_create_get_cancel_mldev",
         parameters=genai_types._CreateBatchJobParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_FLASH,
             src=[{
                 "contents": [{
                     "parts": [{"text": "Why is the sky blue?"}],

google/genai/tests/shared/caches/test_create_get_delete.py

@@ -17,6 +17,7 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 def create_get_delete(client, parameters):
@@ -46,7 +47,7 @@ def create_get_delete(client, parameters):
     pytest_helper.TestTableItem(
         name="test_create_get_delete",
         parameters=genai_types._CreateCachedContentParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_CACHE,
             config=genai_types.CreateCachedContentConfig(
                 contents=[genai_types.Part.from_uri(
                     file_uri="gs://cloud-samples-data/generative-ai/image/a-man-and-a-dog.png",

google/genai/tests/shared/caches/test_create_update_get.py

@@ -17,6 +17,7 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 def create_update_get(client, parameters):
@@ -50,7 +51,7 @@ def create_update_get(client, parameters):
     pytest_helper.TestTableItem(
         name="test_create_update_get",
         parameters=genai_types._CreateCachedContentParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_CACHE,
             config=genai_types.CreateCachedContentConfig(
                 contents=[genai_types.Part.from_uri(
                     file_uri="gs://cloud-samples-data/generative-ai/image/a-man-and-a-dog.png",

google/genai/tests/shared/chats/test_send_message.py

@@ -15,11 +15,12 @@
 
 """Table tests for chats.send_message()."""
 
+import pydantic
 from ... import pytest_helper
-from pydantic import BaseModel
+from .. import model_names
 
 
-class _SendMessageParameters(BaseModel):
+class _SendMessageParameters(pydantic.BaseModel):
   model: str
   message: str
 
@@ -32,7 +33,7 @@ def send_message(client, parameters):
     pytest_helper.TestTableItem(
         name="test_send_message",
         parameters=_SendMessageParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_FLASH,
             message="Hello",
         ),
     ),

google/genai/tests/shared/chats/test_send_message_stream.py

@@ -15,11 +15,12 @@
 
 """Table tests for chats.send_message_stream()."""
 
+import pydantic
 from ... import pytest_helper
-from pydantic import BaseModel
+from .. import model_names
 
 
-class _SendMessageStreamParameters(BaseModel):
+class _SendMessageStreamParameters(pydantic.BaseModel):
     model: str
     message: str
 
@@ -34,7 +35,7 @@ def send_message_stream(client, parameters):
     pytest_helper.TestTableItem(
         name="test_send_message_stream",
         parameters=_SendMessageStreamParameters(
-            model="gemini-2.5-flash",
+            model=model_names.MODEL_FLASH,
             message="Tell a joke.",
         ),
     ),

google/genai/tests/shared/model_names.py

@@ -0,0 +1,30 @@
+# Copyright 2025 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.
+#
+
+"""Central model name constants for shared tests.
+
+Update the model names here to change them across all shared tests.
+After changing model names, re-record replays with:
+  sh run_tests.sh pytest tests/shared --mode=record
+"""
+
+# The default model used for most shared tests (generate, chat, tokens, etc.)
+MODEL_FLASH = 'gemini-3-flash-preview'
+
+# The default model used for caching tests.
+MODEL_CACHE = 'gemini-3-pro-preview'
+
+# The default model used for tuning tests.
+MODEL_FLASH_FOR_TUNING = 'gemini-2.5-flash'

google/genai/tests/shared/models/test_compute_tokens.py

@@ -18,13 +18,14 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 test_table: list[pytest_helper.TestTableItem] = [
     pytest_helper.TestTableItem(
         name='test_compute_tokens',
         parameters=genai_types._ComputeTokensParameters(
-            model='gemini-2.5-flash',
+            model=model_names.MODEL_FLASH,
             contents='The quick brown fox jumps over the lazy dog.',
         ),
         exception_if_mldev=(

google/genai/tests/shared/models/test_count_tokens.py

@@ -18,13 +18,14 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. model_names import MODEL_FLASH
 
 
 test_table: list[pytest_helper.TestTableItem] = [
     pytest_helper.TestTableItem(
         name='test_count_tokens',
         parameters=genai_types._CountTokensParameters(
-            model='gemini-2.5-flash',
+            model=MODEL_FLASH,
             contents='The quick brown fox jumps over the lazy dog.',
         ),
     ),

google/genai/tests/shared/models/test_generate_content.py

@@ -17,20 +17,21 @@
 
 from .... import types as genai_types
 from ... import pytest_helper
+from .. import model_names
 
 
 test_table: list[pytest_helper.TestTableItem] = [
     pytest_helper.TestTableItem(
         name='test_generate_content',
         parameters=genai_types._GenerateContentParameters(
-            model='gemini-2.5-flash',
+            model=model_names.MODEL_FLASH,
             contents='The quick brown fox jumps over the lazy dog.',
         ),
     ),
     pytest_helper.TestTableItem(
         name='test_generate_content_with_config_schema',
         parameters=genai_types._GenerateContentParameters(
-            model='gemini-2.5-flash',
+            model=model_names.MODEL_FLASH,
             contents='Return a summary of the passage.',
             config={'response_mime_type': 'application/json', 'response_schema': {'type': 'OBJECT', 'properties': {'summary': {'type': 'STRING'}}}},
         ),
@@ -39,7 +40,7 @@
     pytest_helper.TestTableItem(
         name='test_generate_content_with_config_json_schema',
         parameters=genai_types._GenerateContentParameters(
-            model='gemini-2.5-flash',
+            model=model_names.MODEL_FLASH,
             contents='Return a JSON summary.',
             config={'response_mime_type': 'application/json', 'response_json_schema': {'type': 'object', 'properties': {'summary': {'type': 'string'}}}},
         ),