Automating Test Suite Generation: AI Agents for Robust, Bug-Free Code

# pip install openai crewai python-dotenv
import os
from crewai import Agent, Task, Process, Crew
from dotenv import load_dotenv

# Load environment variables from .env file
load_dotenv()

# Configure OpenAI API key and model
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")
os.environ["OPENAI_MODEL_NAME"] = "gpt-4o"  # Using GPT-4o for its strong code generation capabilities

code_analyzer = Agent(
    role='Code Analyzer',
    goal='Analyze given Python code to understand its functionality, identify testable units, and describe expected behaviors.',
    backstory='An expert in static code analysis and software design, capable of breaking down complex code into manageable, testable components. Always focuses on identifying inputs, outputs, side effects, and implicit conditions.',
    verbose=True,
    allow_delegation=False
)

strategy_planner = Agent(
    role='Test Strategy Planner',
    goal='Based on the code analysis, devise a comprehensive test strategy. Include specific unit test scenarios, identify crucial edge cases (e.g., boundary conditions, invalid inputs), and detail what needs to be asserted.',
    backstory='A seasoned QA architect specializing in designing robust testing methodologies for critical applications. Adept at anticipating failure points and ensuring maximum coverage.',
    verbose=True,
    allow_delegation=False
)

test_generator = Agent(
    role='Python Pytest Generator',
    goal='Generate production-ready Pytest code snippets based on the provided code and detailed test plan. Focus on correctness, readability, and adherence to Pytest best practices, including proper assertions and parameterization where appropriate.',
    backstory='A senior Python developer and testing expert, proficient in Pytest, mocking, and writing clean, maintainable, and effective test cases that validate all aspects of functionality.',
    verbose=True,
    allow_delegation=False # For this example, we keep it simple without delegation
)

# Example function to be tested
python_code_to_test = """
def calculate_discount(price: float, discount_percentage: float) -> float:
    """Calculates the final price after applying a discount.

    Args:
        price (float): The original price.
        discount_percentage (float): The discount percentage (e.g., 10 for 10%).

    Returns:
        float: The price after discount, or original price if discount is invalid.
    """
    if not (0 <= discount_percentage <= 100):
        return price
    return price * (1 - discount_percentage / 100)
"""

# Tasks
analysis_task = Task(
    description=f"""Analyze the following Python code:
{python_code_to_test}
Identify its purpose, inputs, outputs, and any special conditions or edge cases. Pay close attention to the discount percentage validation logic.
""",
    agent=code_analyzer,
    expected_output='A detailed summary of the code\'s functionality, identified functions, their parameters, return types, and potential areas for testing, including specific conditions like invalid discount ranges.'
)

strategy_task = Task(
    description="""Based on the code analysis, propose a comprehensive test plan for the 'calculate_discount' function. Include specific unit test scenarios, clearly identify critical edge cases like negative prices, zero discount, 100% discount, discounts outside the 0-100 range, and zero price. For each scenario, state the expected input and the precise expected output/behavior.
""",
    agent=strategy_planner,
    context=[analysis_task],
    expected_output='A structured test plan outlining unit test scenarios, specific edge cases with inputs and expected outcomes, and a clear list of assertions needed for the calculate_discount function.'
)

generation_task = Task(
    description=f"""Generate production-ready Pytest code for the following Python function:
{python_code_to_test}

Use the provided test plan to create robust test cases. Ensure the generated code includes a variety of tests covering normal operation, zero values, negative values, and especially boundary conditions (0% and 100% discount) and out-of-range discounts (e.g., -5%, 110%).
Ensure the generated code is production-ready, clean, follows Pytest best practices (e.g., using pytest.approx for floats), and is enclosed in a single Python file format, complete with necessary imports and test functions. Do not include any explanations, just the pure Pytest code.
""",
    agent=test_generator,
    context=[strategy_task],
    expected_output='Production-ready Pytest code in a single Python file format, complete with imports and test functions, specifically for the calculate_discount function.'
)

# Create the Crew and kick off the process
testing_crew = Crew(
    agents=[code_analyzer, strategy_planner, test_generator],
    tasks=[analysis_task, strategy_task, generation_task],
    process=Process.sequential, # Tasks execute in order
    verbose=True # See detailed agent execution logs
)

print("\n--- Kicking off the AI Test Generation Crew ---\n")
result = testing_crew.kickoff()

print("\n\n--- Generated Pytest Code ---\n")
print(result)

import pytest

# Assuming the calculate_discount function is in a file named 'my_module.py'
# For this example, we'll define it directly for demonstration.
def calculate_discount(price: float, discount_percentage: float) -> float:
    """Calculates the final price after applying a discount.

    Args:
        price (float): The original price.
        discount_percentage (float): The discount percentage (e.g., 10 for 10%).

    Returns:
        float: The price after discount, or original price if discount is invalid.
    """
    if not (0 <= discount_percentage <= 100):
        return price
    return price * (1 - discount_percentage / 100)

def test_calculate_discount_normal_case():
    assert calculate_discount(100, 10) == 90.0
    assert calculate_discount(50, 20) == 40.0
    assert calculate_discount(150.0, 5.5) == pytest.approx(141.75)

def test_calculate_discount_zero_discount():
    assert calculate_discount(100, 0) == 100.0

def test_calculate_discount_full_discount():
    assert calculate_discount(100, 100) == 0.0

def test_calculate_discount_invalid_negative_discount():
    assert calculate_discount(100, -5) == 100.0
    assert calculate_discount(0, -10) == 0.0

def test_calculate_discount_invalid_over_hundred_discount():
    assert calculate_discount(100, 110) == 100.0
    assert calculate_discount(50, 200) == 50.0

def test_calculate_discount_zero_price():
    assert calculate_discount(0, 10) == 0.0
    assert calculate_discount(0, 99) == 0.0

def test_calculate_discount_float_values_price():
    assert calculate_discount(99.99, 10) == pytest.approx(89.991)
    assert calculate_discount(1.0, 33.33) == pytest.approx(0.6667)

def test_calculate_discount_large_price_and_discount():
    assert calculate_discount(1000000, 75) == 250000.0

def test_calculate_discount_boundary_values_with_floats():
    assert calculate_discount(100.0, 0.0) == 100.0
    assert calculate_discount(100.0, 100.0) == 0.0

# .github/workflows/ai-test-gen.yml
name: AI Companion Test Generation
on:
  pull_request:
    types: [opened]

jobs:
  generate-tests:
    runs-on: ubuntu-latest
    steps:
      - uses: actions/checkout@v4
      - uses: actions/setup-python@v5
        with:
          python-version: '3.12'

      - name: Generate Missing Tests
        env:
          ANTHROPIC_API_KEY: ${{ secrets.ANTHROPIC_API_KEY }}
        run: |
          python scripts/ai_test_generator.py --pr-diff
          pytest tests/generated/