ada-assistant / src /agents /unit_test_generator.py
syum-af's picture
added unit test generator agent
6ef5a61
Raw
History Blame Contribute Delete
4.16 kB
"""
Unit Test Generator Agent for Ada Conversion Assistant
This agent is responsible for:
1. Analyzing Python code files
2. Reverse-engineering unit tests from existing code
3. Generating comprehensive test suites with pytest
4. Creating test cases that cover edge cases and business logic
"""
from typing import Any, Union
from pathlib import Path
from smolagents import ToolCallingAgent, FinalAnswerTool
from .utils.agent_utils import extract_response_text
class UnitTestGeneratorAgent:
"""Agent specialized in generating unit tests for Python code"""
def __init__(self, model: Any):
"""Initialize agent with LLM model"""
self.model = model
self.code_agent = ToolCallingAgent(
model=self.model,
tools=[FinalAnswerTool()],
)
def generate_unit_tests(self, python_file_path: Union[Path, str]) -> str:
"""
Generate comprehensive unit tests for Python code using LLM
Args:
python_file_path: Path to Python file to generate tests for
Returns:
Generated unit test code as string
"""
# Convert to Path object if string
file_path = Path(python_file_path)
# Read Python file contents
try:
with open(file_path, 'r', encoding='utf-8') as f:
file_content = f.read()
except Exception as e:
return self._create_fallback_test_with_error(f"File read error: {str(e)}")
# Create unit test generation prompt
test_prompt = self._create_test_generation_prompt(file_path, file_content)
try:
# Use the code agent to generate unit tests
result = self.code_agent.run(test_prompt)
# Extract response text from RunResult
response_text = extract_response_text(result)
return response_text
except Exception as e:
return self._create_fallback_test_with_error(f"LLM test generation failed: {str(e)}")
def _create_test_generation_prompt(self, file_path: Path, file_content: str) -> str:
"""Create structured unit test generation prompt for LLM"""
return f"""
Analyze the Python code and generate comprehensive unit tests using pytest.
Create tests that thoroughly cover the functionality, edge cases, and potential error conditions.
Python File: {file_path}
Python Code Content:
{file_content}
Generate unit tests following these guidelines:
1. Use pytest framework with appropriate fixtures
2. Test all public methods and functions
3. Include positive test cases (happy path)
4. Include negative test cases (error conditions)
5. Test edge cases and boundary conditions
6. Use descriptive test method names that explain what is being tested
7. Include proper assertions and expected behavior
8. Mock external dependencies if needed
9. Follow Python testing best practices
10. Ensure tests are isolated and independent
Structure your response as clean Python code with:
- Proper imports (pytest, unittest.mock if needed, etc.)
- Test class(es) with descriptive names
- Test methods with clear, descriptive names
- Proper assertions and error handling tests
- Comments explaining complex test scenarios
Return ONLY the complete Python test code, ready to run with pytest.
"""
def _create_fallback_test_with_error(self, error_msg: str) -> str:
"""Create fallback test response when generation fails"""
return f"""# Unit Test Generation Error
Error: {error_msg}
# Manual test creation recommended
# Please review the Python code manually and create tests using pytest framework.
import pytest
class TestGenerationFailed:
def test_placeholder(self):
# Replace with actual tests for your code
assert True, "Test generation failed - implement manually"
"""