| 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: |
| |
| return "# TODO: Implement integration tests\n" |
|
|
