File size: 2,051 Bytes
bd91486 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | import ast
from typing import List, Dict, Any
import black
class TestGenerator:
def __init__(self):
self.test_templates = {
'python': {
'unit': self._generate_python_unit_test,
'integration': self._generate_python_integration_test
}
}
def generate_tests(self, code: str, language: str, test_type: str = 'unit') -> str:
generator = self.test_templates.get(language, {}).get(test_type)
if not generator:
raise ValueError(f"Unsupported language or test type: {language}/{test_type}")
test_code = generator(code)
try:
return black.format_str(test_code, mode=black.FileMode())
except:
return test_code
def _generate_python_unit_test(self, code: str) -> str:
tree = ast.parse(code)
test_cases = []
for node in ast.walk(tree):
if isinstance(node, ast.FunctionDef):
test_cases.append(self._generate_function_test(node))
return self._format_test_class(test_cases)
def _generate_function_test(self, func_node: ast.FunctionDef) -> str:
args = [arg.arg for arg in func_node.args.args]
test_name = f"test_{func_node.name}"
test_template = f"""
def {test_name}(self):
# Arrange
{''.join(f'{arg} = None # TODO: Add test value\n ' for arg in args)}
# Act
result = {func_node.name}({', '.join(args)})
# Assert
self.assertIsNotNone(result) # TODO: Add specific assertions
"""
return test_template
def _format_test_class(self, test_cases: List[str]) -> str:
return f"""import unittest
class TestGeneratedCode(unittest.TestCase):
{''.join(test_cases)}
if __name__ == '__main__':
unittest.main()
"""
def _generate_python_integration_test(self, code: str) -> str:
# Similar to unit test but with more complex scenarios
return "# TODO: Implement integration tests\n"
|