testgenie-py 0.3.7__py3-none-any.whl → 0.3.8__py3-none-any.whl

testgen/models/db_test_case.py

@@ -0,0 +1,65 @@
+class DBTestCase:
+    def __init__(self, expected_output, inputs, test_function: str, last_run_time, test_method_type: int, test_suite_id: int, function_id: int):
+        self._expected_output = expected_output
+        self._inputs = inputs
+        self._test_function = test_function
+        self._last_run_time = last_run_time
+        self._test_method_type = test_method_type
+        self._test_suite_id = test_suite_id
+        self._function_id = function_id
+
+    @property
+    def expected_output(self):
+        return self._expected_output
+
+    @expected_output.setter
+    def expected_output(self, value):
+        self._expected_output = value
+
+    @property
+    def inputs(self):
+        return self._inputs
+
+    @inputs.setter
+    def inputs(self, value):
+        self._inputs = value
+
+    @property
+    def test_function(self) -> str:
+        return self._test_function
+
+    @test_function.setter
+    def test_function(self, value: str) -> None:
+        self._test_function = value
+
+    @property
+    def last_run_time(self):
+        return self._last_run_time
+
+    @last_run_time.setter
+    def last_run_time(self, value) -> None:
+        self._last_run_time = value
+
+    @property
+    def test_method_type(self) -> int:
+        return self._test_method_type
+
+    @test_method_type.setter
+    def test_method_type(self, value: int) -> None:
+        self._test_method_type = value
+
+    @property
+    def test_suite_id(self) -> int:
+        return self._test_suite_id
+
+    @test_suite_id.setter
+    def test_suite_id(self, value: int) -> None:
+        self._test_suite_id = value
+
+    @property
+    def function_id(self) -> int:
+        return self._function_id
+
+    @function_id.setter
+    def function_id(self, value: int) -> None:
+        self._function_id = value

testgen/models/function.py

@@ -0,0 +1,56 @@
+class Function:
+    def __init__(self, name: str, params, start_line: int, end_line: int, num_lines: int, source_file_id: int):
+        self._name = name
+        self._params = params
+        self._start_line = start_line
+        self._end_line = end_line
+        self._num_lines = num_lines
+        self._source_file_id = source_file_id
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @name.setter
+    def name(self, value: str) -> None:
+        self._name = value
+
+    @property
+    def params(self) -> str:
+        return self._params
+
+    @params.setter
+    def params(self, value: str) -> None:
+        self._params = value
+
+    @property
+    def start_line(self) -> int:
+        return self._start_line
+
+    @start_line.setter
+    def start_line(self, value: int) -> None:
+        self._start_line = value
+
+    @property
+    def end_line(self) -> int:
+        return self._end_line
+
+    @end_line.setter
+    def end_line(self, value: int) -> None:
+        self._end_line = value
+
+    @property
+    def num_lines(self) -> int:
+        return self._num_lines
+
+    @num_lines.setter
+    def num_lines(self, value: int) -> None:
+        self._num_lines = value
+
+    @property
+    def source_file_id(self) -> int:
+        return self._source_file_id
+
+    @source_file_id.setter
+    def source_file_id(self, value: int) -> None:
+        self._source_file_id = value

testgen/models/function_metadata.py

@@ -1,12 +1,14 @@
 import ast
 from types import ModuleType
+from typing import Any
 
 
 class FunctionMetadata:
-    def __init__(self, filename: str, module: ModuleType, class_name: str, function_name: str, func_def: ast.FunctionDef, params: dict):
+    def __init__(self, filename: str, module: ModuleType, class_name: str, func: Any, function_name: str, func_def: ast.FunctionDef, params: dict):
         self._filename: str = filename
         self._module: ModuleType = module
         self._class_name: str = class_name
+        self._func: Any = func
         self._function_name: str = function_name
         self._func_def: ast.FunctionDef = func_def
         self._params: dict = params
@@ -34,6 +36,14 @@ class FunctionMetadata:
     @class_name.setter
     def class_name(self, class_name: str):
         self._class_name = class_name
+
+    @property
+    def func(self) -> Any:
+        return self._func
+
+    @func.setter
+    def func(self, func: Any):
+        self._func = func
     
     @property
     def function_name(self) -> str:

testgen/models/generator_context.py

@@ -1,18 +1,24 @@
 from types import ModuleType
 from typing import List
 
+from pyasn1_modules.rfc6031 import id_ct_KP_sKeyPackage
+
 from testgen.models.test_case import TestCase
 
 
 class GeneratorContext:
-    def __init__(self, filepath: str, filename: str, class_name:str | None, module: ModuleType, output_path: str, test_cases: List[TestCase]):
+    def __init__(self, filepath: str, filename: str, class_name:str | None, module: ModuleType, output_path: str,
+                 test_cases: List[TestCase], is_package: bool, package_name: str, import_path: str):
         self._filepath: str = filepath
         self._filename: str = filename
         self._class_name: str = class_name
         self._module: ModuleType = module
         self._output_path: str = output_path
         self._test_cases: List[TestCase] = test_cases
-    
+        self._is_package: bool = is_package
+        self._package_name: str = package_name
+        self._import_path: str = import_path
+
     @property
     def filepath(self) -> str:
         return self._filepath
@@ -61,3 +67,27 @@ class GeneratorContext:
     def test_cases(self, value: List[TestCase]) -> None:
         self._test_cases = value
 
+    @property
+    def is_package(self) -> bool:
+        return self._is_package
+
+    @is_package.setter
+    def is_package(self, value: bool) -> None:
+        self._is_package = value
+
+    @property
+    def package_name(self) -> str:
+        return self._package_name
+
+    @package_name.setter
+    def package_name(self, value: str) -> None:
+        self._package_name = value
+
+    @property
+    def import_path(self) -> str:
+        return self._import_path
+
+    @import_path.setter
+    def import_path(self, value: str) -> None:
+        self._import_path = value
+

testgen/models/source_file.py

@@ -0,0 +1,29 @@
+class SourceFile:
+    def __init__(self, path: str, lines_of_code: int, last_modified):
+        self._path = path
+        self._lines_of_code = lines_of_code
+        self._last_modified = last_modified
+
+    @property
+    def path(self) -> str:
+        return self._path
+
+    @path.setter
+    def path(self, value: str) -> None:
+        self._path = value
+
+    @property
+    def lines_of_code(self) -> int:
+        return self._lines_of_code
+
+    @lines_of_code.setter
+    def lines_of_code(self, value: int) -> None:
+        self._lines_of_code = value
+
+    @property
+    def last_modified(self):
+        return self._last_modified
+
+    @last_modified.setter
+    def last_modified(self, value) -> None:
+        self._last_modified = value

testgen/models/test_result.py

@@ -0,0 +1,38 @@
+class TestResult:
+    def __init__(self, test_case_id: int, status: bool, error: str, execution_time):
+        self._test_case_id = test_case_id
+        self._status = status
+        self._error = error
+        self._execution_time = execution_time
+
+    @property
+    def test_case_id(self) -> int:
+        return self._test_case_id
+
+    @test_case_id.setter
+    def test_case_id(self, value: int) -> None:
+        self._test_case_id = value
+
+    @property
+    def status(self) -> bool:
+        return self._status
+
+    @status.setter
+    def status(self, value: bool) -> None:
+        self._status = value
+
+    @property
+    def error(self) -> str:
+        return self._error
+
+    @error.setter
+    def error(self, value: str) -> None:
+        self._error = value
+
+    @property
+    def execution_time(self):
+        return self._execution_time
+
+    @execution_time.setter
+    def execution_time(self, value) -> None:
+        self._execution_time = value

testgen/models/test_suite.py

@@ -0,0 +1,20 @@
+class TestSuite:
+    def __init__(self, name: str, creation_date):
+        self._name = name
+        self._creation_date = creation_date
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @name.setter
+    def name(self, value: str) -> None:
+        self._name = value
+
+    @property
+    def creation_date(self):
+        return self._creation_date
+
+    @creation_date.setter
+    def creation_date(self, value) -> None:
+        self._creation_date = value

testgen/reinforcement/agent.py

@@ -128,30 +128,4 @@ class ReinforcementAgent:
         print(f"UPDATING Q TABLE FOR STATE: {state}, ACTION: {action} WITH REWARD: {reward}")
         new_q = (1 - self.learning_rate) * current_q + self.learning_rate * (reward + max_next_q)
         
-        self.q_table[(state, action)] = new_q
-        
-        """def optimize_test_suit(self, current_state, executable_statements):
-        # Try to optimize test cases by repeatedly performing remove actions if reached full coverage
-        test_case_count = current_state[1]
-        optimization_attempts = min(10, test_case_count - 1)
-
-        for _ in range(optimization_attempts):
-            if test_case_count <= 1:
-                break
-
-            action = "remove"
-            next_state, reward = self.env.step(action)
-
-            new_covered = next_state[0]
-            new_uncovered = [stmt for stmt in executable_statements if stmt not in new_covered]
-
-            if len(new_uncovered) == 0:
-                current_state = next_state
-                test_case_count = current_state[2]
-                print(f"Optimized to {test_case_count} test cases.")
-            else:
-                # Add a test case back if removing broke coverage
-                self.env.step("add")
-                break
-
-        return current_state"""
+        self.q_table[(state, action)] = new_q

testgen/reinforcement/environment.py

@@ -4,6 +4,7 @@ from typing import List, Tuple
 
 import coverage
 
+from testgen.models.function_metadata import FunctionMetadata
 from testgen.service.logging_service import get_logger
 import testgen.util.coverage_utils
 import testgen.util.file_utils
@@ -13,11 +14,9 @@ from testgen.models.test_case import TestCase
 
 
 class ReinforcementEnvironment:
-    def __init__(self, file_name, fut: ast.FunctionDef, module, class_name: str | None, initial_test_cases: List[TestCase], state: AbstractState):
-        self.file_name = file_name
-        self.fut = fut
-        self.module = module
-        self.class_name  = class_name
+    def __init__(self, filepath: str, function_data: FunctionMetadata, initial_test_cases: List[TestCase], state: AbstractState):
+        self.filepath = filepath
+        self.function_data = function_data
         self.initial_test_cases = initial_test_cases
         self.test_cases = initial_test_cases.copy()
         self.state = state
@@ -35,13 +34,13 @@ class ReinforcementEnvironment:
 
         # Execute action
         if action == "add":
-            self.test_cases.append(randomizer.new_random_test_case(self.file_name, self.class_name, self.fut))
+            self.test_cases.append(randomizer.new_random_test_case(self.filepath, self.function_data.class_name, self.function_data.func_def))
         elif action == "merge" and len(self.test_cases) > 1:
-            self.test_cases.append(randomizer.combine_cases(self.module, self.test_cases))
+            self.test_cases.append(randomizer.combine_cases(self.function_data.module, self.test_cases))
         elif action == "remove" and len(self.test_cases) > 1:
             self.test_cases = randomizer.remove_case(self.test_cases)
         elif action == "z3":
-            self.test_cases = randomizer.get_z3_test_cases(self.file_name, self.class_name, self.fut, self.test_cases)
+            self.test_cases = randomizer.get_z3_test_cases(self.filepath, self.function_data.class_name, self.function_data.func_def, self.test_cases)
         else:
             raise ValueError("Invalid action")
 
@@ -91,81 +90,6 @@ class ReinforcementEnvironment:
 
         print(f"Final reward {reward}")
         return reward
-
-    
-    def get_all_executable_statements(self):
-        """Get all executable statements including else branches"""
-        import ast
-
-        test_cases = [tc for tc in self.test_cases if tc.func_name == self.fut.name]
-
-        executable_lines = set()
-        if not test_cases:
-            self.logger.debug("Warning: No test cases available to determine executable statements")
-            from testgen.util.randomizer import new_random_test_case
-            temp_case = new_random_test_case(self.file_name, self.class_name, self.fut)
-            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.file_name, self.class_name, self.fut.name, temp_case.inputs)
-            executable_lines.update(analysis[1])  # Add executable lines from coverage analysis
-        else:
-            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.file_name, self.class_name, self.fut.name, test_cases[0].inputs)
-        
-        executable_lines.update(analysis[1])  # Add executable lines from coverage analysis
-        # Get standard executable lines from coverage.py
-        executable_lines = list(executable_lines)
-        
-        # Parse the source file to find else branches
-        with open(self.file_name, 'r') as f:
-            source = f.read()
-        
-        # Parse the code
-        tree = ast.parse(source)        
-        # Find our specific function
-        for node in ast.walk(tree):
-            if isinstance(node, ast.ClassDef) and node.name == self.class_name:
-                # If we have a class, find the method
-                for method in node.body:
-                    if isinstance(method, ast.FunctionDef) and method.name == self.fut.name:
-                        # Find all if statements in this method
-                        for if_node in ast.walk(method):
-                            if isinstance(if_node, ast.If) and if_node.orelse:
-                                # There's an else branch
-                                if isinstance(if_node.orelse[0], ast.If):
-                                    # This is an elif - already counted
-                                    continue
-                                
-                                # Get the line number of the first statement in the else block
-                                # and subtract 1 to get the 'else:' line
-                                else_line = if_node.orelse[0].lineno - 1
-                                
-                                # Check if this is actually an else line (not a nested if)
-                                with open(self.file_name, 'r') as f:
-                                    lines = f.readlines()
-                                    if else_line <= len(lines):
-                                        line_content = lines[else_line - 1].strip()
-                                        if line_content == "else:":
-                                            if else_line not in executable_lines:
-                                                executable_lines.append(else_line)
-            if isinstance(node, ast.FunctionDef) and node.name == self.fut.name:
-                # Find all if statements in this function
-                for if_node in ast.walk(node):
-                    if isinstance(if_node, ast.If) and if_node.orelse:
-                        # There's an else branch
-                        if isinstance(if_node.orelse[0], ast.If):
-                            # This is an elif - already counted
-                            continue
-                        
-                        # Get the line number of the first statement in the else block
-                        # and subtract 1 to get the 'else:' line
-                        else_line = if_node.orelse[0].lineno - 1
-                        
-                        # Check if this is actually an else line (not a nested if)
-                        with open(self.file_name, 'r') as f:
-                            lines = f.readlines()
-                            if else_line <= len(lines):
-                                line_content = lines[else_line - 1].strip()
-                                if line_content == "else:":
-                                    if else_line not in executable_lines:
-                                        executable_lines.append(else_line)
         
         return sorted(executable_lines)
     
@@ -180,13 +104,7 @@ class ReinforcementEnvironment:
         # Execute all test cases
         for test_case in self.test_cases:
             try:
-                module = testgen.util.file_utils.load_module(self.file_name)
-                if self.class_name:
-                    class_obj = getattr(module, self.class_name)
-                    instance = class_obj()
-                    func = getattr(instance, self.fut.name)
-                else:
-                    func = getattr(module, self.fut.name)
+                func = self.function_data.func
                 _ = func(*test_case.inputs)
             except Exception as e:
                 import traceback
@@ -195,7 +113,7 @@ class ReinforcementEnvironment:
         self.cov.stop()
         
         # Get detailed coverage data including branches
-        file_path = os.path.abspath(self.file_name)
+        file_path = os.path.abspath(self.filepath)
         data = self.cov.get_data()
         
         # Extract function-specific coverage
@@ -225,13 +143,13 @@ class ReinforcementEnvironment:
         import ast
         
         try:
-            with open(self.file_name, 'r') as f:
+            with open(self.filepath, 'r') as f:
                 source = f.read()
             
             tree = ast.parse(source)
             
             for node in ast.walk(tree):
-                if isinstance(node, ast.FunctionDef) and node.name == self.fut.name:
+                if isinstance(node, ast.FunctionDef) and node.name == self.function_data.function_name:
                     # Find the first line of the function
                     start_line = node.lineno
                     

testgen/reinforcement/statement_coverage_state.py

@@ -14,11 +14,11 @@ class StatementCoverageState(AbstractState):
         """Returns calculated coverage and length of test cases in a tuple"""
         all_covered_statements = set()
         for test_case in self.environment.test_cases:
-            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.environment.file_name, self.environment.class_name, self.environment.fut.name, test_case.inputs)
+            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.environment.filepath, self.environment.function_data, test_case.inputs)
             covered = testgen.util.coverage_utils.get_list_of_covered_statements(analysis)
             all_covered_statements.update(covered)
 
-        executable_statements = self.environment.get_all_executable_statements()
+        executable_statements = testgen.util.coverage_utils.get_all_executable_statements(self.environment.filepath, self.environment.function_data, self.environment.test_cases)
 
         if not executable_statements or executable_statements == 0:
             calc_coverage = 0.0
@@ -26,10 +26,11 @@ class StatementCoverageState(AbstractState):
             calc_coverage: float = (len(all_covered_statements) / len(executable_statements)) * 100
 
         self.logger.debug(f"GET STATE ALL COVERED STATEMENTS: {all_covered_statements}")
-        self.logger.debug(f"GET STATE ALL EXECUTABLE STATEMENTS: {self.environment.get_all_executable_statements()}")
+        self.logger.debug(f"GET STATE ALL EXECUTABLE STATEMENTS: {executable_statements}")
         self.logger.debug(f"GET STATE FLOAT COVERAGE: {calc_coverage}")
 
         if calc_coverage >= 100:
-            print(f"!!!!!!!!FULLY COVERED FUNCTION: {self.environment.fut.name}!!!!!!!!")
+            print(f"!!!!!!!!FULLY COVERED FUNCTION: {self.environment.function_data.function_name}!!!!!!!!")
+
         return calc_coverage, len(self.environment.test_cases)
             

testgen/service/analysis_service.py

@@ -2,9 +2,10 @@ import inspect
 import ast
 import time
 from types import ModuleType
-from typing import Dict, List
+from typing import Dict, List, Any
 
 import testgen
+from testgen.analyzer.reinforcement_analyzer import ReinforcementAnalyzer
 from testgen.service.logging_service import get_logger
 import testgen.util.file_utils
 import testgen.util.file_utils as file_utils
@@ -37,11 +38,8 @@ class AnalysisService:
 
     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
+            self.test_case_analyzer_context.do_strategy(30)
+            return self.test_case_analyzer_context.test_cases
     
     def create_analysis_context(self, filepath: str) -> AnalysisContext:
         """Create an analysis context for the given file."""
@@ -88,6 +86,7 @@ class AnalysisService:
         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)
+        list_of_function_data: List[FunctionMetadata] = self.get_function_data(filepath, module, class_name)
         functions: List[ast.FunctionDef] = testgen.util.utils.get_functions(tree)
         self.class_name = class_name
         time_limit: int = 30
@@ -95,14 +94,14 @@ class AnalysisService:
 
         q_table = self._load_q_table()
 
-        for function in functions:
-            print(f"\nStarting reinforcement learning for function {function.name}")
+        for function in list_of_function_data:
+            print(f"\nStarting reinforcement learning for function {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, self.class_name, function_test_cases, state=StatementCoverageState(None))
+            environment = ReinforcementEnvironment(filepath, function, function_test_cases, state=StatementCoverageState(None))
             environment.state = StatementCoverageState(environment)
 
             # Create agent with existing Q-table
@@ -115,10 +114,10 @@ class AnalysisService:
                 new_test_cases = agent.collect_test_cases()
                 function_test_cases.extend(new_test_cases)
 
-            print(f"\nNumber of test cases for {function.name}: {len(function_test_cases)}")
+            print(f"\nNumber of test cases for {function.function_name}: {len(function_test_cases)}")
 
             current_coverage: float = environment.run_tests()
-            print(f"Current coverage: {function.name}: {current_coverage}")
+            print(f"Current coverage: {function.function_name}: {current_coverage}")
 
             q_table.update(agent.q_table)
 
@@ -134,8 +133,8 @@ class AnalysisService:
                     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)}")
+            print(f"Final coverage for {function.function_name}: {best_coverage}%")
+            print(f"Final test cases for {function.function_name}: {len(unique_test_cases)}")
 
         self._save_q_table(q_table)
 
@@ -146,12 +145,23 @@ class AnalysisService:
     def _create_function_metadata(self, filename: str, module: ModuleType, class_name: str | None,
                                   func_node: ast.FunctionDef) -> FunctionMetadata:
         function_name = func_node.name
-
+        func = self._get_func_attr(function_name, module, class_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]:
+        return FunctionMetadata(filename, module, class_name, func, function_name, func_node, param_types)
+
+    @staticmethod
+    def _get_func_attr(function_name: str, module: ModuleType, class_name: str | None) -> Any:
+        if class_name:
+            cls = getattr(module, class_name)
+            instance = cls()
+            func = getattr(instance, function_name)
+        else:
+            func = getattr(module, function_name)
+        return func
+
+    @staticmethod
+    def _get_params(func_node: ast.FunctionDef) -> Dict[str, str]:
         # Extract parameter types
         param_types = {}
         for arg in func_node.args.args:
@@ -173,18 +183,17 @@ class AnalysisService:
 
         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
+            analyzer = ReinforcementAnalyzer(analysis_context, mode=self.reinforcement_mode)
         else:
             raise NotImplementedError(f"Test strategy {strategy} not implemented")
-        
+
+        self.test_case_analyzer_context = TestCaseAnalyzerContext(analysis_context, analyzer)
+
     def set_file_path(self, path: str):
         """Set the file path for analysis."""
         self.file_path = path

testgen/service/cfg_service.py

@@ -1,5 +1,7 @@
 import os
 from typing import List
+
+from testgen.models.function_metadata import FunctionMetadata
 from testgen.models.test_case import TestCase
 from testgen.service.logging_service import get_logger
 from testgen.util.coverage_visualizer import CoverageVisualizer
@@ -46,7 +48,7 @@ class CFGService:
         filename = os.path.basename(file_path).replace('.py', '')
 
         for func in analysis_context.function_data:
-            self.visualizer.get_covered_lines(file_path, analysis_context.class_name, func.func_def, test_cases)
+            self.visualizer.get_covered_lines(file_path, func, test_cases)
 
             base_filename = f"{filename}_{func.function_name}_coverage"
             output_filepath = self.get_versioned_filename(visualization_dir, base_filename)