test: add tests for list[pydantic.BaseModel] support and LiveClient preservation

Adewale-1 · Adewale-1 · commit 7b98873998a3 · 2025-06-01T16:41:25.000-04:00
diff --git a/google/genai/tests/models/test_live_client_and_list_type.py b/google/genai/tests/models/test_live_client_and_list_type.py
@@ -0,0 +1,150 @@
+"""Tests to verify both LiveClient classes and list[pydantic.BaseModel] support."""
+
+import inspect
+from typing import List, Optional
+
+import pytest
+from pydantic import BaseModel, Field
+
+from google import genai
+from google.genai import types
+
+
+@pytest.fixture
+def client():
+    """Return a client that uses the replay_session."""
+    client = genai.Client(api_key="test-api-key")
+    return client
+
+
+def test_live_client_classes_exist():
+    """Verify that LiveClient classes exist and have expected attributes."""
+    # Check that LiveClientMessage exists
+    assert hasattr(types, "LiveClientMessage")
+    assert inspect.isclass(types.LiveClientMessage)
+
+    # Check that LiveClientContent exists
+    assert hasattr(types, "LiveClientContent")
+    assert inspect.isclass(types.LiveClientContent)
+
+    # Check that LiveClientRealtimeInput exists
+    assert hasattr(types, "LiveClientRealtimeInput")
+    assert inspect.isclass(types.LiveClientRealtimeInput)
+
+    # Check that LiveClientSetup exists
+    assert hasattr(types, "LiveClientSetup")
+    assert inspect.isclass(types.LiveClientSetup)
+
+    # Check for Dict versions
+    assert hasattr(types, "LiveClientMessageDict")
+    assert hasattr(types, "LiveClientContentDict")
+    assert hasattr(types, "LiveClientRealtimeInputDict")
+    assert hasattr(types, "LiveClientSetupDict")
+
+
+def test_live_client_message_fields():
+    """Verify that LiveClientMessage has expected fields."""
+    # Get the field details
+    fields = types.LiveClientMessage.__fields__
+
+    # Check for expected fields
+    assert "setup" in fields
+    assert "client_content" in fields
+    assert "realtime_input" in fields
+    assert "tool_response" in fields
+
+
+def test_list_pydantic_in_generate_content_response():
+    """Verify that GenerateContentResponse can handle list[pydantic.BaseModel]."""
+
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: List[str]
+
+    # Create a test response
+    response = types.GenerateContentResponse()
+
+    # Assign a list of pydantic models
+    recipes = [
+        Recipe(
+            recipe_name="Chocolate Chip Cookies",
+            ingredients=["Flour", "Sugar", "Chocolate"],
+        ),
+        Recipe(
+            recipe_name="Oatmeal Cookies", ingredients=["Oats", "Flour", "Brown Sugar"]
+        ),
+    ]
+
+    # This assignment would fail with mypy if the type annotation is incorrect
+    response.parsed = recipes
+
+    # Verify assignment worked properly
+    assert response.parsed is not None
+    assert isinstance(response.parsed, list)
+    assert len(response.parsed) == 2
+    assert all(isinstance(item, Recipe) for item in response.parsed)
+
+
+def test_combined_functionality(client):
+    """Test that combines verification of both LiveClient classes and list[pydantic.BaseModel] support."""
+    # 1. Verify LiveClient classes exist
+    assert hasattr(types, "LiveClientMessage")
+    assert inspect.isclass(types.LiveClientMessage)
+
+    # 2. Test list[pydantic.BaseModel] support in generate_content
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: List[str]
+        instructions: Optional[List[str]] = None
+
+    response = client.models.generate_content(
+        model="gemini-1.5-flash",
+        contents="List 2 simple cookie recipes.",
+        config=types.GenerateContentConfig(
+            response_mime_type="application/json",
+            response_schema=list[Recipe],
+        ),
+    )
+
+    # Verify the parsed field contains a list of Recipe objects
+    assert isinstance(response.parsed, list)
+    assert len(response.parsed) > 0
+    assert all(isinstance(item, Recipe) for item in response.parsed)
+
+    # Access a property to verify the type annotation works correctly
+    recipe = response.parsed[0]
+    assert isinstance(recipe.recipe_name, str)
+    assert isinstance(recipe.ingredients, list)
+
+
+def test_live_connect_config_exists():
+    """Verify that LiveConnectConfig exists and has expected attributes."""
+    # Check that LiveConnectConfig exists
+    assert hasattr(types, "LiveConnectConfig")
+    assert inspect.isclass(types.LiveConnectConfig)
+
+    # Check that LiveConnectConfigDict exists
+    assert hasattr(types, "LiveConnectConfigDict")
+
+    # Get the field details if it's a pydantic model
+    if hasattr(types.LiveConnectConfig, "__fields__"):
+        fields = types.LiveConnectConfig.__fields__
+
+        # Check for expected fields (these might vary based on actual implementation)
+        assert "model" in fields
+
+
+def test_live_client_tool_response():
+    """Verify that LiveClientToolResponse exists and has expected attributes."""
+    # Check that LiveClientToolResponse exists
+    assert hasattr(types, "LiveClientToolResponse")
+    assert inspect.isclass(types.LiveClientToolResponse)
+
+    # Check that LiveClientToolResponseDict exists
+    assert hasattr(types, "LiveClientToolResponseDict")
+
+    # Get the field details
+    fields = types.LiveClientToolResponse.__fields__
+
+    # Check for expected fields (these might vary based on actual implementation)
+    assert "function_response" in fields or "tool_outputs" in fields
diff --git a/google/genai/tests/models/test_parsed_list_mypy.py b/google/genai/tests/models/test_parsed_list_mypy.py
@@ -0,0 +1,165 @@
+"""Tests to verify that mypy correctly handles list[pydantic.BaseModel] in response.parsed."""
+
+from typing import List, cast
+import logging
+
+from pydantic import BaseModel
+
+from google.genai import types
+
+# Configure logging
+logger = logging.getLogger(__name__)
+
+
+def test_mypy_with_list_pydantic():
+    """
+    This test doesn't actually run, but it's meant to be analyzed by mypy.
+
+    The code patterns here would have caused mypy errors before the fix,
+    but now should pass type checking with our enhanced types.
+    """
+
+    # Define a Pydantic model for structured output
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: List[str]
+
+    # Create a mock response (simulating what we'd get from the API)
+    response = types.GenerateContentResponse()
+
+    # Before the fix, this next line would cause a mypy error:
+    # Incompatible types in assignment (expression has type "List[Recipe]",
+    # variable has type "Optional[Union[BaseModel, Dict[Any, Any], Enum]]")
+    #
+    # With our fix adding list[pydantic.BaseModel] to the Union, this is now valid:
+    response.parsed = [
+        Recipe(
+            recipe_name="Chocolate Chip Cookies",
+            ingredients=["Flour", "Sugar", "Chocolate"],
+        ),
+        Recipe(
+            recipe_name="Oatmeal Cookies", ingredients=["Oats", "Flour", "Brown Sugar"]
+        ),
+    ]
+
+    # This pattern would require a type cast before the fix
+    if response.parsed is not None:
+        # Before the fix, accessing response.parsed as a list would cause a mypy error
+        # and require a cast:
+        # parsed_items = cast(list[Recipe], response.parsed)
+
+        # With our fix, we can directly use it as a list without casting:
+        recipes = response.parsed
+
+        # We can iterate over the list without casting
+        for recipe in recipes:
+            logger.info(f"Recipe: {recipe.recipe_name}")
+            for ingredient in recipe.ingredients:
+                logger.info(f" - {ingredient}")
+
+        # We can access elements by index without casting
+        first_recipe = recipes[0]
+        logger.info(f"First recipe: {first_recipe.recipe_name}")
+
+
+def test_with_pydantic_inheritance():
+    """Test with inheritance to ensure the type annotation works with subclasses."""
+
+    class FoodItem(BaseModel):
+        name: str
+
+    class Recipe(FoodItem):
+        ingredients: List[str]
+
+    response = types.GenerateContentResponse()
+
+    # Before the fix, this would require a cast with mypy
+    # Now it works directly with our enhanced type annotation
+    response.parsed = [
+        Recipe(
+            name="Chocolate Chip Cookies", ingredients=["Flour", "Sugar", "Chocolate"]
+        ),
+        Recipe(name="Oatmeal Cookies", ingredients=["Oats", "Flour", "Brown Sugar"]),
+    ]
+
+    if response.parsed is not None:
+        # Previously would need: cast(list[Recipe], response.parsed)
+        recipes = response.parsed
+
+        # Access fields from parent class
+        for recipe in recipes:
+            logger.info(f"Recipe name: {recipe.name}")
+
+
+def test_with_nested_list_models():
+    """Test with nested list models to ensure complex structures work."""
+
+    class Ingredient(BaseModel):
+        name: str
+        amount: str
+
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: List[Ingredient]
+
+    response = types.GenerateContentResponse()
+
+    # With the fix, mypy correctly handles this complex structure
+    response.parsed = [
+        Recipe(
+            recipe_name="Chocolate Chip Cookies",
+            ingredients=[
+                Ingredient(name="Flour", amount="2 cups"),
+                Ingredient(name="Sugar", amount="1 cup"),
+            ],
+        ),
+        Recipe(
+            recipe_name="Oatmeal Cookies",
+            ingredients=[
+                Ingredient(name="Oats", amount="1 cup"),
+                Ingredient(name="Flour", amount="1.5 cups"),
+            ],
+        ),
+    ]
+
+    if response.parsed is not None:
+        recipes = response.parsed
+
+        # Access nested structures without casting
+        for recipe in recipes:
+            logger.info(f"Recipe: {recipe.recipe_name}")
+            for ingredient in recipe.ingredients:
+                logger.info(f" - {ingredient.name}: {ingredient.amount}")
+
+
+# Example of how you would previously need to cast the results
+def old_approach_with_cast():
+    """
+    This demonstrates the old approach that required explicit casting,
+    which was less type-safe and more error-prone.
+    """
+
+    class Recipe(BaseModel):
+        recipe_name: str
+        ingredients: List[str]
+
+    response = types.GenerateContentResponse()
+
+    # Simulate API response
+    response.parsed = [
+        Recipe(
+            recipe_name="Chocolate Chip Cookies",
+            ingredients=["Flour", "Sugar", "Chocolate"],
+        ),
+        Recipe(
+            recipe_name="Oatmeal Cookies", ingredients=["Oats", "Flour", "Brown Sugar"]
+        ),
+    ]
+
+    if response.parsed is not None:
+        # Before our fix, you'd need this cast for mypy to be happy
+        recipes = cast(List[Recipe], response.parsed)
+
+        # Using the cast list
+        for recipe in recipes:
+            logger.info(f"Recipe: {recipe.recipe_name}")
diff --git a/google/genai/tests/models/test_parsed_list_support.py b/google/genai/tests/models/test_parsed_list_support.py
