testgenie-py 0.1.0__py3-none-any.whl

testgen/service/service.py

@@ -0,0 +1,389 @@
+import inspect
+import json
+import os
+import re
+import sqlite3
+import time
+import subprocess
+from types import ModuleType
+from typing import List
+import testgen.util.file_utils as file_utils
+
+from testgen.models.test_case import TestCase
+from testgen.service.analysis_service import AnalysisService
+from testgen.service.generator_service import GeneratorService
+from testgen.sqlite.db_service import DBService
+from testgen.models.analysis_context import AnalysisContext
+from testgen.util.coverage_visualizer import CoverageVisualizer
+
+# Constants for test strategies
+AST_STRAT = 1
+FUZZ_STRAT = 2
+RANDOM_STRAT = 3
+REINFORCE_STRAT = 4
+
+# Constants for test formats
+UNITTEST_FORMAT = 1
+PYTEST_FORMAT = 2
+DOCTEST_FORMAT = 3
+
+class Service:
+    def __init__(self):
+        self.test_strategy: int = 0
+        self.test_format: int = 0
+        self.file_path = None
+        self.generated_file_path = None
+        self.class_name = None
+        self.test_cases = []
+        self.reinforcement_mode = "train"
+        
+        # Initialize specialized services
+        self.analysis_service = AnalysisService()
+        self.generator_service = GeneratorService(None, None, None)
+        self.db_service = DBService()
+
+    def select_all_from_db(self) -> None:
+        rows = self.db_service.get_test_suites()
+        for row in rows:
+            print(repr(dict(row)))
+
+    def generate_tests(self, output_path=None):
+        """Generate tests for a class or module."""
+        module = file_utils.load_module(self.file_path)
+        class_name = self.analysis_service.get_class_name(module)
+
+        if self.test_strategy == AST_STRAT:
+            self.generated_file_path = self.generate_function_code()
+            Service.wait_for_file(self.generated_file_path)
+            self.analysis_service.set_file_path(self.generated_file_path)
+            module = file_utils.load_module(self.generated_file_path)
+            class_name = self.analysis_service.get_class_name(module)
+            self.analysis_service.set_test_strategy(self.test_strategy, module.__name__, self.class_name)
+
+        else:
+            self.analysis_service.set_test_strategy(self.test_strategy, module.__name__, self.class_name)
+
+        test_cases: List[TestCase] = []
+        if self.test_strategy == REINFORCE_STRAT:
+            test_cases = self.analysis_service.do_reinforcement_learning(self.file_path, self.reinforcement_mode)
+        else:
+            test_cases = self.analysis_service.generate_test_cases()
+
+        self.test_cases = test_cases
+
+        file_path_to_use = self.generated_file_path if self.test_strategy == AST_STRAT else self.file_path
+        self.db_service.save_test_generation_data(
+            file_path_to_use,
+            test_cases,
+            self.test_strategy,
+            class_name
+        )
+
+        if os.environ.get("RUNNING_IN_DOCKER") is not None:
+            print(f"Serializing test cases {test_cases}")
+            self.serialize_test_cases(test_cases)
+            return None  # Exit early in analysis-only mode
+
+        test_file = self.generate_test_file(test_cases, output_path, module, class_name)
+
+        # Ensure the test file is ready
+        Service.wait_for_file(test_file)
+        return test_file
+
+    def generate_test_file(self, test_cases: List[TestCase], output_path: str | None = None,
+                           module: ModuleType | None = None, class_name: str | None = None) -> str:
+        if module is None:
+            module = file_utils.load_module(self.file_path)
+            class_name = self.analysis_service.get_class_name(module)
+
+        if self.test_strategy == AST_STRAT:
+            self.generator_service = GeneratorService(self.generated_file_path, output_path, self.test_format)
+            self.generator_service.set_test_format(self.test_format)
+            module = file_utils.load_module(self.generated_file_path)
+        else:
+            # Create the correct instance of the generator service
+            self.generator_service = GeneratorService(self.file_path, output_path, self.test_format)
+            self.generator_service.set_test_format(self.test_format)
+
+        test_file = self.generator_service.generate_test_file(
+            module,
+            class_name,
+            test_cases,
+            output_path
+        )
+
+        # Ensure the test file is ready
+        Service.wait_for_file(test_file)
+        return test_file
+
+    def generate_function_code(self):
+        """Generate function code for a given class or module."""
+        module = file_utils.load_module(self.file_path)
+        class_name = self.analysis_service.get_class_name(module)
+        functions = self.inspect_class(class_name)
+        return self.generator_service.generate_function_code(self.file_path, class_name, functions)
+    
+    def run_coverage(self, test_file):
+        """Run coverage analysis on the generated tests."""
+        Service.wait_for_file(test_file)
+        file_path_to_use = self.generated_file_path if self.test_strategy == AST_STRAT else self.file_path
+        coverage_output = ""
+        
+        try:
+            if self.test_format == UNITTEST_FORMAT:
+                subprocess.run(["python", "-m", "coverage", "run", "--source=.", "-m", "unittest", test_file], check=True)
+                result = subprocess.run(
+                    ["python", "-m", "coverage", "report", file_path_to_use], 
+                    check=True, 
+                    capture_output=True, 
+                    text=True
+                )
+                coverage_output = result.stdout
+            elif self.test_format == PYTEST_FORMAT:
+                subprocess.run(["python", "-m", "coverage", "run", "--source=.", "-m", "pytest", test_file], check=True)
+                result = subprocess.run(
+                    ["python", "-m", "coverage", "report", file_path_to_use], 
+                    check=True,
+                    capture_output=True,
+                    text=True
+                )
+                coverage_output = result.stdout
+            elif self.test_format == DOCTEST_FORMAT:
+                result = subprocess.run(
+                    ["python", "-m", "coverage", "run", "--source=.", "-m", "doctest", "-v", test_file], 
+                    check=True,
+                    capture_output=True,
+                    text=True
+                )
+                coverage_output = result.stdout
+            else:
+                raise ValueError("Unsupported test format for coverage analysis.")
+                
+            self._save_coverage_data(coverage_output, file_path_to_use)
+            
+        except subprocess.CalledProcessError as e:
+            raise RuntimeError(f"Error running coverage subprocess: {e}")
+
+    def _save_coverage_data(self, coverage_output, file_path):
+        """Parse coverage output and save to database."""
+        try:
+            lines = coverage_output.strip().split('\n')
+            for line in lines:
+                if file_path in line:
+                    parts = line.split()
+                    if len(parts) >= 4:
+                        file_name = os.path.basename(file_path)
+                        try:
+                            total_lines = int(parts[-3])
+                            missed_lines = int(parts[-2])
+                            executed_lines = total_lines - missed_lines
+                            coverage_str = parts[-1].strip('%')
+                            branch_coverage = float(coverage_str) / 100
+                            
+                            source_file_id = self._get_source_file_id(file_path)
+                            
+                            self.db_service.insert_coverage_data(
+                                file_name, 
+                                executed_lines, 
+                                missed_lines, 
+                                branch_coverage, 
+                                source_file_id
+                            )
+                            break
+                        except (ValueError, IndexError) as e:
+                            print(f"Error parsing coverage data: {e}")
+        except Exception as e:
+            print(f"Error saving coverage data: {e}")
+            
+    def _get_source_file_id(self, file_path):
+        """Helper to get source file ID from DB."""
+        conn = sqlite3.connect(self.db_service.db_name)
+        conn.row_factory = sqlite3.Row
+        cursor = conn.cursor()
+        
+        cursor.execute("SELECT id FROM SourceFile WHERE path = ?", (file_path,))
+        row = cursor.fetchone()
+        
+        conn.close()
+        
+        if row:
+            return row[0]
+        else:
+            with open(file_path, 'r') as f:
+                lines_of_code = len(f.readlines())
+            return self.db_service.insert_source_file(file_path, lines_of_code)
+
+    def serialize_test_cases(self, test_cases):
+        """Serialize input arguments ot JSON-compatible format"""
+        print("##TEST_CASES_BEGIN##")
+        serialized = []
+        for tc in test_cases:
+            case_data = {
+                "func_name": tc.func_name,
+                "inputs": self.serialize_value(tc.inputs),
+                "expected": self.serialize_value(tc.expected)
+            }
+            serialized.append(case_data)
+        print(json.dumps(serialized))
+        print("##TEST_CASES_END##")
+
+    def serialize_value(self, value):
+        """Serialize a single value to a JSON-compatible format"""
+        if value is None or isinstance(value, (bool, int, float, str)):
+            return value
+        elif isinstance(value, (list, tuple)):
+            return [self.serialize_value(v) for v in value]
+        else:
+            return str(value)
+
+    @staticmethod
+    def parse_test_cases_from_logs(logs_output):
+        """Extract and parse test cases from container logs"""
+        pattern = r"##TEST_CASES_BEGIN##\n(.*?)\n##TEST_CASES_END##"
+        match = re.search(pattern, logs_output, re.DOTALL)
+
+        if not match:
+            raise ValueError("Could not find test cases in the container logs")
+
+        test_cases_json = match.group(1).strip()
+        test_cases_data = json.loads(test_cases_json)
+
+        test_cases: List[TestCase] = []
+        for tc_data in test_cases_data:
+            test_case = TestCase(
+                func_name=tc_data["func_name"],
+                inputs=tc_data["inputs"],
+                expected=tc_data["expected"]
+            )
+            test_cases.append(test_case)
+
+        return test_cases
+
+    def set_file_path(self, path: str):
+        """Set the file path for analysis and validate it."""
+        if os.path.isfile(path) and path.endswith(".py"):
+            self.file_path = path
+            self.analysis_service.set_file_path(path)
+        else:
+            raise ValueError("Invalid file path! Please provide a valid Python file path.")
+
+    def set_class_name(self, class_name: str):
+        """Set the class name to analyze."""
+        self.class_name = class_name
+
+    def set_test_generator_format(self, test_format: int):
+        """Set the test generator format."""
+        self.test_format = test_format
+        self.generator_service.set_test_format(test_format)
+
+    def set_test_analysis_strategy(self, strategy: int):
+        """Set the test analysis strategy."""
+        self.test_strategy = strategy
+        module = file_utils.load_module(self.file_path)
+        self.analysis_service.set_test_strategy(strategy, module.__name__, self.class_name)
+    
+    def get_analysis_context(self, filepath: str) -> AnalysisContext:
+        """Create an analysis context for the given file."""
+        return self.analysis_service.create_analysis_context(filepath)
+
+    @staticmethod
+    def wait_for_file(file_path, retries=5, delay=1):
+        """Wait for the generated file to appear."""
+        while retries > 0 and not os.path.exists(file_path):
+            time.sleep(delay)
+            retries -= 1
+        if not os.path.exists(file_path):
+            raise FileNotFoundError(f"File '{file_path}' not found after waiting.")
+
+    def get_full_import_path(self) -> str:
+        """Get the full import path for the current file."""
+        package_root = self.find_package_root()
+        if not package_root:
+            raise ImportError(f"Could not determine the package root for {self.file_path}.")
+
+        module_path = os.path.abspath(self.file_path)
+        rel_path = os.path.relpath(module_path, package_root)
+        package_path = rel_path.replace(os.sep, ".")
+        
+        if package_path.endswith(".py"):
+            package_path = package_path[:-3]
+
+        package_name = os.path.basename(package_root)
+        if not package_path.startswith(package_name + "."):
+            package_path = package_name + "." + package_path
+
+        return package_path
+
+    def find_package_root(self):
+        """Find the package root directory."""
+        current_dir = os.path.abspath(os.path.dirname(self.file_path))
+        last_valid = None
+
+        while current_dir:
+            if "__init__.py" in os.listdir(current_dir):
+                last_valid = current_dir
+            else:
+                break
+
+            parent_dir = os.path.dirname(current_dir)
+            if parent_dir == current_dir:
+                break
+            current_dir = parent_dir
+
+        return last_valid
+
+    def inspect_class(self, class_name):
+        """Inspect a class or module and return its functions."""
+        module = file_utils.load_module(self.file_path)
+
+        # Handle module-level functions when class_name is None
+        if not class_name:
+            # Get module-level functions
+            functions = inspect.getmembers(module, inspect.isfunction)
+            return functions
+
+        # Handle class functions
+        cls = getattr(module, class_name, None)
+        if cls is None:
+            raise ValueError(f"Class '{class_name}' not found in module '{module.__name__}'.")
+
+        functions = inspect.getmembers(cls, inspect.isfunction)
+        return functions
+
+    @staticmethod
+    def resolve_module_path(module_name):
+        """Resolve a module name to its file path by checking multiple locations."""
+        direct_path = f"/controller/{module_name}.py"
+        if os.path.exists(direct_path):
+            print(f"Found module at {direct_path}")
+            return direct_path
+
+        testgen_path = f"/controller/testgen/{module_name}.py"
+        if os.path.exists(testgen_path):
+            print(f"Found module at {testgen_path}")
+            return testgen_path
+
+        if '.' in module_name:
+            parts = module_name.split('.')
+            potential_path = os.path.join('/controller', *parts) + '.py'
+            if os.path.exists(potential_path):
+                print(f"Found module at {potential_path}")
+                return potential_path
+
+        print(f"Could not find module: {module_name}")
+        return None
+
+    def visualize_test_coverage(self):
+        from testgen.service.cfg_service import CFGService
+        cfg_service = CFGService()
+        cfg_service.initialize_visualizer(self)
+
+        return cfg_service.visualize_test_coverage(
+            file_path=self.file_path,
+            test_cases=self.test_cases,
+        )
+
+    def set_reinforcement_mode(self, mode: str):
+        self.reinforcement_mode = mode
+        if hasattr(self ,'analysis_service'):
+            self.analysis_service.set_reinforcement_mode(mode)

testgen/sqlite/db.py

@@ -0,0 +1,84 @@
+import sqlite3
+
+def create_database(db_name="testgen.db"):
+    conn = sqlite3.connect(db_name)
+    cursor = conn.cursor()
+
+    # Enable foreign key constraints
+    cursor.execute("PRAGMA foreign_keys = ON;")
+
+    # Create tables
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS TestSuite (
+            id INTEGER PRIMARY KEY,
+            name TEXT,
+            creation_date TIMESTAMP
+        );
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS SourceFile (
+            id INTEGER PRIMARY KEY,
+            path TEXT,
+            lines_of_code INTEGER,
+            last_modified TIMESTAMP
+        );
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS Function (
+            id INTEGER PRIMARY KEY,
+            name TEXT,
+            start_line INTEGER,
+            end_line INTEGER,
+            num_lines INTEGER,
+            source_file_id INTEGER,
+            FOREIGN KEY (source_file_id) REFERENCES SourceFile(id)
+        );
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS TestCase (
+            id INTEGER PRIMARY KEY,
+            name TEXT,
+            expected_output TEXT,  -- storing JSON as TEXT
+            input TEXT,            -- storing JSON as TEXT
+            test_function TEXT,
+            last_run_time TIMESTAMP,
+            test_method_type INTEGER,
+            test_suite_id INTEGER,
+            function_id INTEGER,
+            FOREIGN KEY (test_suite_id) REFERENCES TestSuite(id),
+            FOREIGN KEY (function_id) REFERENCES Function(id)
+        );
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS TestResult (
+            id INTEGER PRIMARY KEY,
+            test_case_id INTEGER,
+            status BOOLEAN,
+            error TEXT,
+            execution_time TIMESTAMP,
+            FOREIGN KEY (test_case_id) REFERENCES TestCase(id)
+        );
+    """)
+
+    cursor.execute("""
+        CREATE TABLE IF NOT EXISTS CoverageData (
+            id INTEGER PRIMARY KEY,
+            file_name TEXT,
+            executed_lines INTEGER,
+            missed_lines INTEGER,
+            branch_coverage REAL,
+            source_file_id INTEGER,
+            FOREIGN KEY (source_file_id) REFERENCES SourceFile(id)
+        );
+    """)
+
+    conn.commit()
+    conn.close()
+    print(f"Database '{db_name}' created successfully with all tables.")
+
+if __name__ == "__main__":
+    create_database()

testgen/sqlite/db_service.py

@@ -0,0 +1,219 @@
+import os
+import sqlite3
+import json
+import time
+import ast
+from datetime import datetime
+from typing import List, Tuple
+
+from testgen.models.test_case import TestCase
+from testgen.sqlite.db import create_database
+
+class DBService:
+    def __init__(self, db_name="testgen.db"):
+        """Initialize database service with connection to specified database."""
+        self.db_name = db_name
+        self.conn = None
+        self.cursor = None
+        self._connect()
+
+    def _connect(self):
+        """Establish connection to the database."""
+        if not os.path.exists(self.db_name):
+            create_database(self.db_name)
+        
+        self.conn = sqlite3.connect(self.db_name)
+        self.conn.row_factory = sqlite3.Row
+        self.cursor = self.conn.cursor()
+        # Enable foreign keys
+        self.cursor.execute("PRAGMA foreign_keys = ON;")
+
+    def close(self):
+        """Close the database connection."""
+        if self.conn:
+            self.conn.close()
+            self.conn = None
+            self.cursor = None
+
+    def insert_test_suite(self, name: str) -> int:
+        """Insert a test suite and return its ID."""
+        self.cursor.execute(
+            "INSERT INTO TestSuite (name, creation_date) VALUES (?, ?)",
+            (name, datetime.now())
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def insert_source_file(self, path: str, lines_of_code: int) -> int:
+        """Insert a source file and return its ID."""
+        # Check if file already exists
+        self.cursor.execute("SELECT id FROM SourceFile WHERE path = ?", (path,))
+        existing = self.cursor.fetchone()
+        if existing:
+            return existing[0]
+        
+        self.cursor.execute(
+            "INSERT INTO SourceFile (path, lines_of_code, last_modified) VALUES (?, ?, ?)",
+            (path, lines_of_code, datetime.now())
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def insert_function(self, name: str, start_line: int, end_line: int, source_file_id: int) -> int:
+        """Insert a function and return its ID."""
+        num_lines = end_line - start_line + 1
+        
+        # Check if function already exists for this source file
+        self.cursor.execute(
+            "SELECT id FROM Function WHERE name = ? AND source_file_id = ?", 
+            (name, source_file_id)
+        )
+        existing = self.cursor.fetchone()
+        if existing:
+            return existing[0]
+            
+        self.cursor.execute(
+            "INSERT INTO Function (name, start_line, end_line, num_lines, source_file_id) VALUES (?, ?, ?, ?, ?)",
+            (name, start_line, end_line, num_lines, source_file_id)
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def insert_test_case(self, test_case: TestCase, test_suite_id: int, function_id: int, test_method_type: int) -> int:
+        """Insert a test case and return its ID."""
+        # Convert inputs and expected output to JSON strings
+        inputs_json = json.dumps(test_case.inputs)
+        expected_json = json.dumps(test_case.expected)
+        
+        self.cursor.execute(
+            "INSERT INTO TestCase (name, expected_output, input, test_function, last_run_time, test_method_type, test_suite_id, function_id) "
+            "VALUES (?, ?, ?, ?, ?, ?, ?, ?)",
+            (
+                f"test_{test_case.func_name}",
+                expected_json,
+                inputs_json,
+                test_case.func_name,
+                datetime.now(),
+                test_method_type,
+                test_suite_id,
+                function_id
+            )
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def insert_test_result(self, test_case_id: int, status: bool, error: str = None) -> int:
+        """Insert a test result and return its ID."""
+        self.cursor.execute(
+            "INSERT INTO TestResult (test_case_id, status, error, execution_time) VALUES (?, ?, ?, ?)",
+            (test_case_id, status, error, datetime.now())
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def insert_coverage_data(self, file_name: str, executed_lines: int, missed_lines: int, 
+                           branch_coverage: float, source_file_id: int) -> int:
+        """Insert coverage data and return its ID."""
+        self.cursor.execute(
+            "INSERT INTO CoverageData (file_name, executed_lines, missed_lines, branch_coverage, source_file_id) "
+            "VALUES (?, ?, ?, ?, ?)",
+            (file_name, executed_lines, missed_lines, branch_coverage, source_file_id)
+        )
+        self.conn.commit()
+        return self.cursor.lastrowid
+
+    def save_test_generation_data(self, file_path: str, test_cases: List[TestCase], 
+                                test_method_type: int, class_name: str = None) -> Tuple[int, List[int]]:
+        """
+        Save all data related to a test generation run.
+        Returns the test suite ID and a list of test case IDs.
+        """
+        # Count lines in the source file
+        with open(file_path, 'r') as f:
+            lines_of_code = len(f.readlines())
+        
+        # Create test suite
+        strategy_names = {1: "AST", 2: "Fuzz", 3: "Random", 4: "Reinforcement"}
+        suite_name = f"{strategy_names.get(test_method_type, 'Unknown')}_Suite_{int(time.time())}"
+        test_suite_id = self.insert_test_suite(suite_name)
+        
+        # Insert source file
+        source_file_id = self.insert_source_file(file_path, lines_of_code)
+        
+        # Process functions and test cases
+        test_case_ids = []
+        function_ids = {}  # Cache function IDs to avoid redundant queries
+        
+        for test_case in test_cases:
+            # Extract function name from test case
+            func_name = test_case.func_name
+            
+            # If function not already processed
+            if func_name not in function_ids:
+                # Get function line numbers
+                start_line, end_line = self._get_function_line_numbers(file_path, func_name)
+                function_id = self.insert_function(func_name, start_line, end_line, source_file_id)
+                function_ids[func_name] = function_id
+            
+            # Insert test case
+            test_case_id = self.insert_test_case(
+                test_case, 
+                test_suite_id, 
+                function_ids[func_name], 
+                test_method_type
+            )
+            test_case_ids.append(test_case_id)
+        
+        return test_suite_id, test_case_ids
+    
+    def _get_function_line_numbers(self, file_path: str, function_name: str) -> Tuple[int, int]:
+        """
+        Extract the start and end line numbers for a function in a file.
+        Returns a tuple of (start_line, end_line).
+        """
+        try:
+            # Load the file and parse it
+            with open(file_path, 'r') as f:
+                file_content = f.read()
+                
+            tree = ast.parse(file_content)
+            
+            # Find the function definition
+            for node in ast.walk(tree):
+                if isinstance(node, ast.FunctionDef) and node.name == function_name:
+                    end_line = node.end_lineno if hasattr(node, 'end_lineno') else node.lineno + 5  # Estimate if end_lineno not available
+                    return node.lineno, end_line
+                    
+            # Also look for class methods
+            for node in ast.walk(tree):
+                if isinstance(node, ast.ClassDef):
+                    for class_node in node.body:
+                        if isinstance(class_node, ast.FunctionDef) and class_node.name == function_name:
+                            end_line = class_node.end_lineno if hasattr(class_node, 'end_lineno') else class_node.lineno + 5
+                            return class_node.lineno, end_line
+        except Exception as e:
+            print(f"Error getting function line numbers: {e}")
+            
+        # If we reach here, the function wasn't found or there was an error
+        return 0, 0
+    
+    def get_test_suites(self):
+        """Get all test suites from the database."""
+        self.cursor.execute("SELECT * FROM TestSuite ORDER BY creation_date DESC")
+        return self.cursor.fetchall()
+    
+    def get_test_cases_by_function(self, function_name):
+        """Get all test cases for a specific function."""
+        self.cursor.execute(
+            "SELECT tc.* FROM TestCase tc JOIN Function f ON tc.function_id = f.id WHERE f.name = ?",
+            (function_name,)
+        )
+        return self.cursor.fetchall()
+    
+    def get_coverage_by_file(self, file_path):
+        """Get coverage data for a specific file."""
+        self.cursor.execute(
+            "SELECT cd.* FROM CoverageData cd JOIN SourceFile sf ON cd.source_file_id = sf.id WHERE sf.path = ?", 
+            (file_path,)
+        )
+        return self.cursor.fetchall()