	In Browser
	StumbleUpon
	del.icio.us
	Google
	Google Buzz
	reddit
	LinkedIn

	Facebook
	Twitter
	Linkedin
	E-Mail

Generative AI > Google Gemini API > Test Case Generation Agent

Test Case Generation Agent

Author: Venkata Sudhakar

Writing comprehensive test cases is time-consuming but critical for software quality. A test case generation agent can analyse function signatures and business logic, then produce pytest-compatible tests covering normal behaviour, edge cases, and invalid inputs - dramatically accelerating the QA cycle.

In this tutorial, we build a ShopMax India test generation agent that reads the source of business logic functions and produces ready-to-run pytest suites. The agent covers happy paths, boundary values, and exception scenarios automatically.

The below example shows the agent analysing a discount calculation function and generating a complete test suite.

import os
from google import genai
from google.genai import types

# The business function to be tested
SOURCE_CODE = """
def calculate_discount(price: float, customer_tier: str, quantity: int) -> float:
    Calculate the final discounted price for ShopMax India orders.
    Tiers: "bronze" = 5%, "silver" = 10%, "gold" = 15%, "platinum" = 20%
    Additional 2% for quantity >= 10 units.
    Raises ValueError if price <= 0, quantity <= 0, or unknown tier.
    Returns final price after all discounts applied.
    tier_discounts = {"bronze": 0.05, "silver": 0.10, "gold": 0.15, "platinum": 0.20}
    if price <= 0 or quantity <= 0:
        raise ValueError("Price and quantity must be positive")
    if customer_tier not in tier_discounts:
        raise ValueError(f"Unknown tier: {customer_tier}")
    discount = tier_discounts[customer_tier]
    if quantity >= 10:
        discount += 0.02
    return round(price * (1 - discount), 2)
"""

client = genai.Client(api_key=os.environ["GEMINI_API_KEY"])

prompt = f"""You are a senior Python QA engineer.
Analyse the following function and generate a complete pytest test suite.
Cover: happy path for each tier, boundary at quantity=9 vs 10, negative inputs, unknown tier.
Output ONLY valid Python pytest code, no explanations.

Function to test:
{SOURCE_CODE}
"""

response = client.models.generate_content(
    model="gemini-2.0-flash",
    contents=prompt,
    config=types.GenerateContentConfig(
        response_mime_type="text/plain"
    )
)

generated_tests = response.text
print(generated_tests)

It gives the following output,

import pytest
from discount import calculate_discount

class TestHappyPath:
    def test_bronze_tier(self):
        assert calculate_discount(1000.0, "bronze", 1) == 950.0

    def test_silver_tier(self):
        assert calculate_discount(1000.0, "silver", 1) == 900.0

    def test_gold_tier(self):
        assert calculate_discount(1000.0, "gold", 1) == 850.0

    def test_platinum_tier(self):
        assert calculate_discount(1000.0, "platinum", 1) == 800.0

class TestBulkDiscount:
    def test_quantity_9_no_bulk(self):
        assert calculate_discount(1000.0, "silver", 9) == 900.0

    def test_quantity_10_gets_bulk(self):
        assert calculate_discount(1000.0, "silver", 10) == 880.0

    def test_quantity_50_gets_bulk(self):
        assert calculate_discount(1000.0, "gold", 50) == 830.0

class TestInvalidInputs:
    def test_negative_price(self):
        with pytest.raises(ValueError):
            calculate_discount(-100.0, "silver", 1)

    def test_zero_quantity(self):
        with pytest.raises(ValueError):
            calculate_discount(1000.0, "silver", 0)

    def test_unknown_tier(self):
        with pytest.raises(ValueError, match="Unknown tier"):
            calculate_discount(1000.0, "diamond", 1)

Running pytest gives,

tests/test_discount.py::TestHappyPath::test_bronze_tier PASSED
tests/test_discount.py::TestHappyPath::test_silver_tier PASSED
tests/test_discount.py::TestHappyPath::test_gold_tier PASSED
tests/test_discount.py::TestHappyPath::test_platinum_tier PASSED
tests/test_discount.py::TestBulkDiscount::test_quantity_9_no_bulk PASSED
tests/test_discount.py::TestBulkDiscount::test_quantity_10_gets_bulk PASSED
tests/test_discount.py::TestBulkDiscount::test_quantity_50_gets_bulk PASSED
tests/test_discount.py::TestInvalidInputs::test_negative_price PASSED
tests/test_discount.py::TestInvalidInputs::test_zero_quantity PASSED
tests/test_discount.py::TestInvalidInputs::test_unknown_tier PASSED

10 passed in 0.12s

Integrate this agent into your CI pipeline to auto-generate regression tests whenever a function changes. Pass the git diff of changed functions as input instead of static source code for targeted test generation. For complex class hierarchies, include the class definition and its dependencies in the prompt for more accurate mocks and fixtures.

Send your comments, suggestions or queries regarding this site to [email protected].