testgenie-py 0.1.0__py3-none-any.whl

testgen/service/analysis_service.py

@@ -0,0 +1,260 @@
+import inspect
+import ast
+import time
+from types import ModuleType
+from typing import Dict, List
+
+import testgen
+import testgen.util.file_utils
+import testgen.util.file_utils as file_utils
+import testgen.util.utils
+from testgen.analyzer.ast_analyzer import ASTAnalyzer
+from testgen.analyzer.fuzz_analyzer import FuzzAnalyzer
+from testgen.analyzer.random_feedback_analyzer import RandomFeedbackAnalyzer
+from testgen.models.test_case import TestCase
+from testgen.analyzer.test_case_analyzer_context import TestCaseAnalyzerContext
+from testgen.reinforcement.agent import ReinforcementAgent
+from testgen.reinforcement.environment import ReinforcementEnvironment
+from testgen.reinforcement.statement_coverage_state import StatementCoverageState
+from testgen.models.analysis_context import AnalysisContext
+from testgen.models.function_metadata import FunctionMetadata
+
+# Constants for test strategies
+AST_STRAT = 1
+FUZZ_STRAT = 2
+RANDOM_STRAT = 3
+REINFORCE_STRAT = 4
+
+class AnalysisService:
+    def __init__(self):
+        self.file_path = None
+        self.class_name = None
+        self.test_case_analyzer_context = TestCaseAnalyzerContext(None, None)
+        self.test_strategy = 0
+        self.reinforcement_mode = "train"
+
+    def generate_test_cases(self) -> List[TestCase]:
+            """Generate test cases using the current strategy."""
+            if self.test_strategy == REINFORCE_STRAT:
+                return self.do_reinforcement_learning(self.file_path)
+            else:
+                self.test_case_analyzer_context.do_logic()
+                return self.test_case_analyzer_context.test_cases
+    
+    def create_analysis_context(self, filepath: str) -> AnalysisContext:
+        """Create an analysis context for the given file."""
+        filename = file_utils.get_filename(filepath)
+        module = file_utils.load_module(filepath)
+        class_name = self.get_class_name(module)
+        function_data = self.get_function_data(filename, module, class_name)
+        return AnalysisContext(filepath, filename, class_name, module, function_data)
+    
+    def get_function_data(self, filename: str, module: ModuleType, class_name: str | None) -> List[FunctionMetadata]:
+        function_metadata_list: List[FunctionMetadata] = []
+
+        # Parse the module's source code
+        source_code = inspect.getsource(module)
+        tree = ast.parse(source_code)
+
+        # Find all function definitions in the module
+        for node in tree.body:
+            if isinstance(node, ast.FunctionDef):
+                function_metadata_list.append(self._create_function_metadata(filename, module, class_name, node))
+            elif isinstance(node, ast.ClassDef):
+                # Get the class name to check if it matches the expected class
+                ast_class_name = node.name
+
+                # If we have a specified class and it doesn't match, skip
+                if self.class_name and ast_class_name != self.class_name:
+                    continue
+
+                # Process class methods
+                for class_node in node.body:
+                    if isinstance(class_node, ast.FunctionDef):
+                        # Skip private methods (those starting with _)
+                        if not class_node.name.startswith('_'):
+                            function_metadata_list.append(self._create_function_metadata(filename, module, class_name, class_node))
+
+        return function_metadata_list
+
+    def do_reinforcement_learning(self, filepath: str, mode: str = None) -> List[TestCase]:
+        mode = mode or self.reinforcement_mode
+        module: ModuleType = testgen.util.file_utils.load_module(filepath)
+        tree: ast.Module = testgen.util.file_utils.load_and_parse_file_for_tree(filepath)
+        functions: List[ast.FunctionDef] = testgen.util.utils.get_functions(tree)
+        time_limit: int = 30
+        all_test_cases: List[TestCase] = []
+
+        q_table = self._load_q_table()
+
+        for function in functions:
+            print(f"\nStarting reinforcement learning for function {function.name}")
+            start_time = time.time()
+            function_test_cases: List[TestCase] = []
+            best_coverage: float = 0.0
+
+            # Create environment and agent once per function
+            environment = ReinforcementEnvironment(filepath, function, module, function_test_cases, state=StatementCoverageState(None))
+            environment.state = StatementCoverageState(environment)
+
+            # Create agent with existing Q-table
+            agent = ReinforcementAgent(filepath, environment, function_test_cases, q_table)
+
+            if mode == "train":
+                new_test_cases = agent.do_q_learning()
+                function_test_cases.extend(new_test_cases)
+            else:
+                new_test_cases = agent.collect_test_cases()
+                function_test_cases.extend(new_test_cases)
+
+            print(f"\nTest cases for {function.name}: {len(function_test_cases)}")
+
+            current_coverage: float = environment.run_tests()
+            print(f"Current coverage: {function.name}: {current_coverage}")
+
+            q_table.update(agent.q_table)
+
+            # Process and filter unique test cases
+            seen = set()
+            unique_test_cases: List[TestCase] = []
+            for case in function_test_cases:
+                # Make case tuple hashable
+                case_inputs = tuple(case.inputs) if isinstance(case.inputs, list) else case.inputs
+                case_key = (case.func_name, case_inputs)
+                if case_key not in seen:
+                    seen.add(case_key)
+                    unique_test_cases.append(case)
+
+            all_test_cases.extend(unique_test_cases)
+            print(f"Final coverage for {function.name}: {best_coverage}%")
+            print(f"Final test cases for {function.name}: {len(unique_test_cases)}")
+
+        self._save_q_table(q_table)
+
+        print("\nReinforcement Learning Complete")
+        print(f"Total test cases found: {len(all_test_cases)}")
+        return all_test_cases
+
+    def _create_function_metadata(self, filename: str, module: ModuleType, class_name: str | None,
+                                  func_node: ast.FunctionDef) -> FunctionMetadata:
+        function_name = func_node.name
+
+        param_types = self._get_params(func_node)
+
+        return FunctionMetadata(filename, module, class_name, function_name, func_node, param_types)
+    
+    def _get_params(self, func_node: ast.FunctionDef) -> Dict[str, str]:
+        # Extract parameter types
+        param_types = {}
+        for arg in func_node.args.args:
+            param_name = arg.arg
+            if param_name == 'self':
+                continue
+
+            if arg.annotation:
+                param_type = ast.unparse(arg.annotation)
+                param_types[param_name] = param_type
+            else:
+                param_types[param_name] = None
+        return param_types
+    
+    def set_test_strategy(self, strategy: int, module_name: str, class_name: str):
+        """Set the test analysis strategy."""
+        self.test_strategy = strategy
+        analysis_context = self.create_analysis_context(self.file_path)
+
+        if strategy == AST_STRAT:
+            analyzer = ASTAnalyzer(analysis_context)
+            self.test_case_analyzer_context = TestCaseAnalyzerContext(analysis_context, analyzer)
+        elif strategy == FUZZ_STRAT:
+            analyzer = FuzzAnalyzer(analysis_context)
+            self.test_case_analyzer_context = TestCaseAnalyzerContext(analysis_context, analyzer)
+        elif strategy == RANDOM_STRAT:
+            analyzer = RandomFeedbackAnalyzer(analysis_context)
+            self.test_case_analyzer_context = TestCaseAnalyzerContext(analysis_context, analyzer)
+        elif strategy == REINFORCE_STRAT:
+            pass
+        else:
+            raise NotImplementedError(f"Test strategy {strategy} not implemented")
+        
+    def set_file_path(self, path: str):
+        """Set the file path for analysis."""
+        self.file_path = path
+
+    def set_reinforcement_mode(self, mode: str):
+        self.reinforcement_mode = mode
+
+    @staticmethod
+    def get_class_name(module: ModuleType) -> str | None:
+        """Get class from module otherwise return None."""
+        for name, cls in inspect.getmembers(module, inspect.isclass):
+            if cls.__module__ == module.__name__:
+                return name
+        return None
+
+    @staticmethod
+    def _save_q_table(q_table):
+        """Save Q-table to a global JSON file"""
+        import json
+        import os
+
+        q_table_dir = "q_table"
+        os.makedirs(q_table_dir, exist_ok=True)
+        q_table_path = os.path.join(q_table_dir, "global_q_table.json")
+
+        global_q_table = {}
+        if os.path.exists(q_table_path):
+            try:
+                with open(q_table_path, 'r') as f:
+                    global_q_table = json.load(f)
+            except Exception as e:
+                print(f"Error loading existing Q-table: {e}")
+
+        serializable_q_table = {}
+        for key, value in q_table.items():
+            state, action = key
+            state_str = str(state)
+            serializable_q_table[f"{state_str}|{action}"] = value
+
+        global_q_table = serializable_q_table  # Replace with latest Q-table
+
+        try:
+            with open(q_table_path, 'w') as f:
+                json.dump(global_q_table, f)
+            print(f"Q-table saved to {q_table_path}")
+        except Exception as e:
+            print(f"Error saving Q-table: {e}")
+
+    @staticmethod
+    def _load_q_table():
+        """Load Q-table from the global JSON file"""
+        import json
+        import os
+        import ast
+
+        q_table_dir = "q_table"
+        q_table_path = os.path.join(q_table_dir, "global_q_table.json")
+
+        if not os.path.exists(q_table_path):
+            print(f"No existing Q-table found at {q_table_path}")
+            return {}
+
+        try:
+            with open(q_table_path, 'r') as f:
+                serialized_q_table = json.load(f)
+
+            # Convert serialized keys back to (state, action) tuples
+            q_table = {}
+            for key, value in serialized_q_table.items():
+                state_str, action = key.split('|')
+                try:
+                    state = ast.literal_eval(state_str)
+                    q_table[(state, action)] = value
+                except (ValueError, SyntaxError):
+                    print(f"Skipping invalid state: {state_str}")
+
+            print(f"Loaded global Q-table with {len(q_table)} entries")
+            return q_table
+        except Exception as e:
+            print(f"Error loading Q-table: {e}")
+            return {}

testgen/service/cfg_service.py

@@ -0,0 +1,55 @@
+import os
+from typing import List
+from testgen.models.test_case import TestCase
+from testgen.util.coverage_visualizer import CoverageVisualizer
+from testgen.service.analysis_service import AnalysisService
+
+
+class CFGService:
+    """Service for generating and managing Control Flow Graph visualizations."""
+    def __init__(self):
+        self.analysis_service = AnalysisService()
+        self.visualizer = None
+
+    def initialize_visualizer(self, service):
+        self.visualizer = CoverageVisualizer()
+        self.visualizer.set_service(service)
+
+    @staticmethod
+    def create_visualization_directory() -> str:
+        """Create visualization directory if it doesn't exist."""
+        visualization_dir = os.path.join(os.getcwd(), "visualize")
+        if not os.path.exists(visualization_dir):
+            os.makedirs(visualization_dir)
+            print(f"Created visualization directory: {visualization_dir}")
+        return visualization_dir
+
+    @staticmethod
+    def get_versioned_filename(directory: str, base_filename: str) -> str:
+        """Generate a versioned filename to avoid overwriting existing files."""
+        version = 1
+        output_path = os.path.join(directory, f"{base_filename}.png")
+
+        while os.path.exists(output_path):
+            output_path = os.path.join(directory, f"{base_filename}_v{version}.png")
+            version += 1
+
+        return output_path
+
+    def visualize_test_coverage(self, file_path: str, test_cases: List[TestCase]) -> str | None:
+        visualization_dir = self.create_visualization_directory()
+
+        analysis_context = self.analysis_service.create_analysis_context(file_path)
+
+        filename = os.path.basename(file_path).replace('.py', '')
+
+        for func in analysis_context.function_data:
+            self.visualizer.get_covered_lines(file_path, func.func_def, test_cases)
+
+            base_filename = f"{filename}_{func.function_name}_coverage"
+            output_filepath = self.get_versioned_filename(visualization_dir, base_filename)
+
+            self.visualizer.generate_colored_cfg(func.function_name, output_filepath)
+
+        print(f"Generated CFG visualizations in {visualization_dir}")
+        return visualization_dir

testgen/service/generator_service.py

@@ -0,0 +1,169 @@
+import os
+from types import ModuleType
+from typing import List
+
+from testgen.models.test_case import TestCase
+from testgen.generator.code_generator import CodeGenerator
+from testgen.generator.doctest_generator import DocTestGenerator
+from testgen.generator.pytest_generator import PyTestGenerator
+from testgen.generator.unit_test_generator import UnitTestGenerator
+from testgen.inspector.inspector import Inspector
+from testgen.tree.node import Node
+from testgen.tree.tree_utils import build_binary_tree
+from testgen.models.generator_context import GeneratorContext
+
+# Constants for test formats
+UNITTEST_FORMAT = 1
+PYTEST_FORMAT = 2
+DOCTEST_FORMAT = 3
+
+class GeneratorService:
+    def __init__(self, filepath: str, output_path: str, test_format: int = UNITTEST_FORMAT):
+        self.filepath = filepath
+        self.output_path = output_path
+        self.test_format = test_format
+        self.code_generator = CodeGenerator()
+        self.test_generator = UnitTestGenerator(generator_context=None)
+        self.generated_file_path = None
+
+    def set_test_format(self, test_format: int):
+        """Set the test generator format."""
+        self.test_format = test_format
+        if test_format == UNITTEST_FORMAT:
+            print("SETTING TEST FORMAT TO UNITTEST")
+            self.test_generator = UnitTestGenerator(generator_context=None)
+        elif test_format == PYTEST_FORMAT:
+            print("SETTING TEST FORMAT TO PYTEST")
+            self.test_generator = PyTestGenerator(generator_context=None)
+        elif test_format == DOCTEST_FORMAT:
+            print("SETTING TEST FORMAT TO DOCTEST")
+            self.test_generator = DocTestGenerator(generator_context=None)
+        else:
+            raise NotImplementedError(f"Test format {test_format} not implemented")
+
+    def generate_test_file(self, module: ModuleType, class_name: str | None, test_cases: List[TestCase], output_path=None) -> str:
+        """Generate a test file for the given test cases."""
+
+        filename = self.get_filename(self.filepath)
+
+        output_path = self.get_test_file_path(module.__name__, output_path)
+        
+        context = GeneratorContext(
+            filepath=self.filepath,
+            filename=filename,
+            class_name=class_name,
+            module=module,
+            output_path=output_path,
+            test_cases=test_cases
+        )
+
+        self.test_generator.generator_context = context
+        
+        self.test_generator.generate_test_header()
+
+        print("GENERATE TEST FILE: Generated test header")
+
+        self.generate_function_tests(test_cases)
+
+        print("GENERATE TEST FILE: Generate function tests")
+
+        print()
+
+        if self.test_format == DOCTEST_FORMAT:
+            print("SAVING DOCT TEST FILE")
+            self.test_generator.save_file()
+            return self.filepath
+        else:
+            self.test_generator.save_file()
+            return output_path
+
+    def generate_function_tests(self, test_cases: List[TestCase]) -> None:
+        """Generate test functions for the given test cases."""
+        for i, test_case in enumerate(test_cases):
+            unique_func_name = f"{test_case.func_name}_{i}"
+            cases = [(test_case.inputs, test_case.expected)]
+            self.test_generator.generate_test_function(unique_func_name, test_case.func_name, cases)
+
+    def get_test_file_path(self, module_name: str, specified_path=None) -> str:
+        """Determine the path for the generated test file."""
+        if specified_path is not None:
+            if os.path.exists(specified_path):
+                if os.path.isdir(specified_path):
+                    self.ensure_init_py(specified_path)
+                    return os.path.join(specified_path, f"test_{module_name.lower()}.py")
+                else:
+                    print(f"Specified directory path: {specified_path} is not a directory.")
+            else:
+                print(f"Specified directory path: {specified_path} does not exist.")
+
+        current_dir = os.getcwd()
+        test_dir = os.path.join(current_dir, "tests")
+
+        if os.path.exists(test_dir):
+            if os.path.isdir(test_dir):
+                self.ensure_init_py(test_dir)
+                return os.path.join(test_dir, f"test_{module_name.lower()}.py")
+            else:
+                print(f"Test directory path: {test_dir} is not a directory.")
+        else:
+            print(f"Test directory path: {test_dir} does not exist, creating it.")
+            os.mkdir(test_dir)
+            self.ensure_init_py(test_dir)
+
+        return os.path.join(test_dir, f"test_{module_name.lower()}.py")
+
+    def ensure_init_py(self, directory: str):
+        """Ensures that an __init__.py file exists in the given directory."""
+        init_file = os.path.join(directory, "__init__.py")
+        if not os.path.exists(init_file):
+            with open(init_file, "w") as f:
+                pass  # Create an empty __init__.py file
+            print(f"Created __init__.py in {directory}")
+
+    def generate_function_code(self, file_path: str, class_name: str | None, functions: list) -> str:
+        """Generate function code for a given class and its functions."""
+        trees = self.build_func_trees(functions)
+        if class_name:
+            print(class_name)
+            self.generated_file_path = f"generated_{class_name.lower()}.py"
+        else:
+            print(self.get_filename(file_path))
+            self.generated_file_path = f"generated_{self.get_filename(file_path)}"
+
+        # Create the class file
+        if class_name:
+            file = self.code_generator.generate_class(class_name)
+            file.close()
+        else:
+            file = open(self.generated_file_path, "w")
+            file.close()
+
+        # Append function implementations
+        with open(self.generated_file_path, "a") as file:
+            for func, root, params in trees:
+                is_class_method = True if class_name else False
+                code = self.code_generator.generate_code_from_tree(func.__name__, root, params, func, is_class_method)
+                if class_name:
+                    for line in code.split("\n"):
+                        file.write(f"    {line}\n")
+                else:
+                    file.write(code)
+                file.write("\n")
+
+        return self.generated_file_path
+
+    def build_func_trees(self, functions: list):
+        """Build binary trees for function signatures."""
+        tree_list = []
+        for name, func in functions:
+            signature = Inspector.get_signature(func)
+            params = Inspector.get_params_not_self(signature)
+            root = Node(None)
+            build_binary_tree(root, 0, len(params))
+            tree_list.append((func, root, params))
+        return tree_list
+
+    @staticmethod
+    def get_filename(filepath: str) -> str:
+        """Get filename from filepath."""
+        return os.path.basename(filepath)