testgenie-py 0.1.6__py3-none-any.whl → 0.1.8__py3-none-any.whl

testgen/analyzer/random_feedback_analyzer.py

@@ -113,7 +113,7 @@ class RandomFeedbackAnalyzer(TestCaseAnalyzerStrategy, ABC):
             self.covered_lines[func.function_name] = set()
 
         for test_case in [tc for tc in self.test_cases if tc.func_name == func.function_name]:
-            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath,
+            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath, self._analysis_context.class_name,
                                                                          func.function_name, test_case.inputs)
             covered = coverage_utils.get_list_of_covered_statements(analysis)
             self.covered_lines[func.function_name].update(covered)
@@ -249,10 +249,10 @@ class RandomFeedbackAnalyzer(TestCaseAnalyzerStrategy, ABC):
             print("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._analysis_context.filepath, func.func_def)
-            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath, func.function_name,
+            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath, self._analysis_context.class_name, func.function_name,
                                                                          temp_case.inputs)
         else:
-            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath, func.function_name, test_cases[0].inputs)
+            analysis = coverage_utils.get_coverage_analysis(self._analysis_context.filepath, self._analysis_context.class_name, func.function_name, test_cases[0].inputs)
 
         # Get standard executable lines from coverage.py
         executable_lines = list(analysis[1])

testgen/controller/cli_controller.py

@@ -6,7 +6,8 @@ import sys
 import docker
 from docker import DockerClient
 from docker import errors
-
+from testgen.service.logging_service import LoggingService, get_logger
+from testgen.util.file_utils import adjust_file_path_for_docker, get_project_root_in_docker
 from testgen.controller.docker_controller import DockerController
 from testgen.service.service import Service
 from testgen.presentation.cli_view import CLIView
@@ -28,6 +29,85 @@ class CLIController:
         self.view = view
 
     def run(self):
+
+        parser = self.add_arguments()
+
+        args = parser.parse_args()
+
+        LoggingService.get_instance().initialize(
+            debug_mode=args.debug if hasattr(args, 'debug') else False,
+            log_file=args.log_file if hasattr(args, 'log_file') else None,
+            console_output=True
+        )
+
+        logger = get_logger()
+
+        if args.select_all:
+            self.view.display_message("Selecting all from SQLite database...")
+            # Assuming you have a method in your service to handle this
+            self.service.select_all_from_db()
+            return
+            
+        running_in_docker = os.environ.get("RUNNING_IN_DOCKER") is not None
+        if running_in_docker:
+            args.file_path = adjust_file_path_for_docker(args.file_path)
+            self.execute_generation(args, True)
+        elif args.safe and not running_in_docker:
+            client = self.docker_available()
+            # Skip Docker-dependent operations if client is None
+            if client is None and args.safe:
+                self.view.display_message("Running with --safe flag requires Docker. Continuing without safe mode.")
+                args.safe = False
+                self.execute_generation(args)
+            else:
+                docker_controller = DockerController()
+                project_root = get_project_root_in_docker(args.file_path)
+                successful: bool = docker_controller.run_in_docker(project_root, client, args)
+                if not successful:
+                    if hasattr(args, 'db') and args.db:
+                        self.service.db_service = DBService(args.db)
+                        self.view.display_message(f"Using database: {args.db}")
+                    self.execute_generation(args)
+                # Else successful, do nothing - we're done
+        else:
+            # Initialize database service with specified path
+            if hasattr(args, 'db') and args.db:
+                self.service.db_service = DBService(args.db)
+                self.view.display_message(f"Using database: {args.db}")
+            self.view.display_message("Running in local mode...")
+            self.execute_generation(args)
+
+    def execute_generation(self, args: argparse.Namespace, running_in_docker: bool = False):
+        try:
+            self.set_service_args(args)
+
+            if running_in_docker:
+                self.view.display_message("Running in Docker mode...")
+                self.service.generate_test_cases()
+
+            else:
+                test_file = self.service.generate_tests(args.output)
+                self.view.display_message(f"Unit tests saved to: {test_file}")
+                self.view.display_message("Running coverage...")
+                self.service.run_coverage(test_file)
+                self.view.display_message("Tests and coverage data saved to database.")
+
+                if args.visualize:
+                    self.service.visualize_test_coverage()
+
+        except Exception as e:
+            self.view.display_error(f"An error occurred: {e}")
+            # Make sure to close the DB connection on error
+            if hasattr(self.service, 'db_service'):
+                self.service.db_service.close()
+
+    def set_service_args(self, args: argparse.Namespace):
+        self.service.set_file_path(args.file_path)
+        self.service.set_debug_mode(args.debug)
+        self.set_test_format(args)
+        self.set_test_strategy(args)
+
+    def add_arguments(self) -> argparse.ArgumentParser:
         parser = argparse.ArgumentParser(description="A CLI tool for generating unit tests.")
         parser.add_argument("file_path", type=str, help="Path to the Python file.")
         parser.add_argument("--output", "-o", type=str, help="Path to output directory.")
@@ -75,108 +155,48 @@ class CLIController:
             action="store_true",
             help = "Visualize the tests with graphviz"
         )
+        parser.add_argument(
+            "--debug",
+            action="store_true",
+            help="Enable debug logging"
+        )
+        parser.add_argument(
+            "--log-file",
+            type=str,
+            help="Path to log file (if not specified, logs will only go to console)"
+        )
+        return parser
 
-        args = parser.parse_args()
-
-        if args.select_all:
-            self.view.display_message("Selecting all from SQLite database...")
-            # Assuming you have a method in your service to handle this
-            self.service.select_all_from_db()
-            return
-        
-        # Initialize database service with specified path
-        if hasattr(args, 'db') and args.db:
-            self.service.db_service = DBService(args.db)
-            self.view.display_message(f"Using database: {args.db}")
-            
-        running_in_docker = os.environ.get("RUNNING_IN_DOCKER") is not None
-        if running_in_docker:
-            args.file_path = self.adjust_file_path_for_docker(args.file_path)
-            self.execute_generation(args)
-        elif args.safe and not running_in_docker:
-            client = self.docker_available()
-            # Skip Docker-dependent operations if client is None
-            if client is None and args.safe:
-                self.view.display_message("Running with --safe flag requires Docker. Continuing without safe mode.")
-                args.safe = False
-            docker_controller = DockerController()
-            project_root = self.get_project_root_in_docker(args.file_path)
-            successful: bool = docker_controller.run_in_docker(project_root, client, args)
-            if not successful:
-                self.execute_generation(args)
+    def set_test_format(self, args: argparse.Namespace):
+        if args.test_format == "pytest":
+            self.service.set_test_generator_format(PYTEST_FORMAT)
+        elif args.test_format == "doctest":
+            self.service.set_test_generator_format(DOCTEST_FORMAT)
         else:
-            self.view.display_message("Running in local mode...")
-            self.execute_generation(args)
+            self.service.set_test_generator_format(UNITTEST_FORMAT)
 
-    def execute_generation(self, args: argparse.Namespace):
-        try:
-            self.service.set_file_path(args.file_path)
-            if args.test_format == "pytest":
-                self.service.set_test_generator_format(PYTEST_FORMAT)
-            elif args.test_format == "doctest":
-                self.service.set_test_generator_format(DOCTEST_FORMAT)
-            else:
-                self.service.set_test_generator_format(UNITTEST_FORMAT)
-            if args.test_mode == "random":
-                self.view.display_message("Using Random Feedback-Directed Test Generation Strategy.")
-                self.service.set_test_analysis_strategy(RANDOM_STRAT)
-            elif args.test_mode == "fuzz":
-                self.view.display_message("Using Fuzz Test Generation Strategy...")
-                self.service.set_test_analysis_strategy(FUZZ_STRAT)
-            elif args.test_mode == "reinforce":
-                self.view.display_message("Using Reinforcement Learning Test Generation Strategy...")
-                if args.reinforce_mode == "train":
-                    self.view.display_message("Training mode enabled - will update Q-table")
-                else:
-                    self.view.display_message("Training mode disabled - will use existing Q-table")
-                self.service.set_test_analysis_strategy(REINFORCE_STRAT)
-                self.service.set_reinforcement_mode(args.reinforce_mode)
+    def set_test_strategy(self, args: argparse.Namespace):
+        if args.test_mode == "random":
+            self.view.display_message("Using Random Feedback-Directed Test Generation Strategy.")
+            self.service.set_test_analysis_strategy(RANDOM_STRAT)
+        elif args.test_mode == "fuzz":
+            self.view.display_message("Using Fuzz Test Generation Strategy...")
+            self.service.set_test_analysis_strategy(FUZZ_STRAT)
+        elif args.test_mode == "reinforce":
+            self.view.display_message("Using Reinforcement Learning Test Generation Strategy...")
+            if args.reinforce_mode == "train":
+                self.view.display_message("Training mode enabled - will update Q-table")
             else:
-                self.view.display_message("Generating function code using AST analysis...")
-                generated_file_path = self.service.generate_function_code()
-                self.view.display_message(f"Generated code saved to: {generated_file_path}")
-                if not args.generate_only:
-                    self.view.display_message("Using Simple AST Traversal Test Generation Strategy...")
-                    self.service.set_test_analysis_strategy(AST_STRAT)
-
-            test_file = self.service.generate_tests(args.output)
-            self.view.display_message(f"Unit tests saved to: {test_file}")
-            self.view.display_message("Running coverage...")
-            self.service.run_coverage(test_file)
-            self.view.display_message("Tests and coverage data saved to database.")
-
-            if args.visualize:
-                self.service.visualize_test_coverage()
-
-        except Exception as e:
-            self.view.display_error(f"An error occurred: {e}")
-            # Make sure to close the DB connection on error
-            if hasattr(self.service, 'db_service'):
-                self.service.db_service.close()
-
-    def adjust_file_path_for_docker(self, file_path) -> str:
-        file_dir = os.path.abspath(os.path.dirname(file_path))
-        sys.path.append(file_dir)
-        sys.path.append('/controller')
-        file_abs_path = os.path.abspath(file_path)
-        if not os.path.exists(file_abs_path):
-            testgen_path = os.path.join('/controller/testgen', os.path.basename(file_path))
-            if os.path.exists(testgen_path):
-                file_path = testgen_path
-            else:
-                app_path = os.path.join('/controller', os.path.basename(file_path))
-                if os.path.exists(app_path):
-                    file_path = app_path
-        return file_path
-
-    def get_project_root_in_docker(self, script_path) -> str:
-        script_path = os.path.abspath(sys.argv[0])
-        print(f"Script path: {script_path}")
-        script_dir = os.path.dirname(script_path)
-        print(f"Script directory: {script_dir}")
-        project_root = os.path.dirname(script_dir)
-        print(f"Project root directory: {project_root}")
-        return project_root
+                self.view.display_message("Training mode disabled - will use existing Q-table")
+            self.service.set_test_analysis_strategy(REINFORCE_STRAT)
+            self.service.set_reinforcement_mode(args.reinforce_mode)
+        else:
+            self.view.display_message("Generating function code using AST analysis...")
+            generated_file_path = self.service.generate_function_code()
+            self.view.display_message(f"Generated code saved to: {generated_file_path}")
+            if not args.generate_only:
+                self.view.display_message("Using Simple AST Traversal Test Generation Strategy...")
+                self.service.set_test_analysis_strategy(AST_STRAT)
 
     def docker_available(self) -> DockerClient | None:
         try:

testgen/controller/docker_controller.py

@@ -6,6 +6,7 @@ from docker import DockerClient, client
 from docker import errors
 from docker.models.containers import Container
 
+from testgen.service.logging_service import get_logger
 from testgen.service.service import Service
 
 AST_STRAT = 1
@@ -19,10 +20,13 @@ DOCTEST_FORMAT = 3
 class DockerController:
     def __init__(self):
         self.service = Service()
+        self.debug_mode = False
         self.args = None
+        self.logger = get_logger()
 
     def run_in_docker(self, project_root: str, docker_client: DockerClient, args: Namespace) -> bool:
         self.args = args
+        self.debug_mode = True if args.debug else False
         os.environ["RUNNING_IN_DOCKER"] = "1"
 
         # Check if Docker image exists, build it if not
@@ -30,18 +34,19 @@ class DockerController:
         # If args.safe is set to false it means the image was not found and the system will try to run_locally
         self.get_image(docker_client, image_name, project_root)
         if not self.args.safe:
-            print("Docker image not found. Running locally...")
+            self.logger.info("Docker image not found. Running locally...")
             return False
 
-        docker_args = [os.path.basename(args.file_path)] + [arg for arg in sys.argv[2:] if arg != "--safe"]
+        docker_args = [args.file_path] + [arg for arg in sys.argv[2:] if arg != "--safe"]
 
         # Run the container with the same arguments
         try:
+            self.debug(f"project_root: {project_root}")
             container = self.run_container(docker_client, image_name, docker_args, project_root)
 
             # Stream the logs to the console
             logs_output = self.get_logs(container)
-            print(logs_output)
+            self.debug(logs_output)
 
             try:
                 # Create the target directory if it doesn't exist
@@ -51,14 +56,14 @@ class DockerController:
                     target_path = args.output
                 os.makedirs(target_path, exist_ok=True)
 
-                print(f"SERVICE target path after logs: {target_path}")
+                self.debug(f"SERVICE target path after logs: {target_path}")
 
                 test_cases = self.service.parse_test_cases_from_logs(logs_output)
 
                 print(f"Extracted {len(test_cases)} test cases from container.")
 
                 file_path = os.path.abspath(args.file_path)
-                print(f"Filepath in CLI CONTROLLER: {file_path}")
+                self.debug(f"Filepath in CLI CONTROLLER: {file_path}")
                 self.service.set_file_path(file_path)
 
                 if args.test_format == "pytest":
@@ -73,9 +78,10 @@ class DockerController:
 
                 if not args.generate_only:
                     print("Running coverage...")
-                    import traceback
-                    print(traceback.format_exc())
                     self.service.run_coverage(test_file)
+                    
+                # Add explicit return True here
+                return True
 
             except Exception as e:
                 print(f"Error running container: {e}")
@@ -99,7 +105,7 @@ class DockerController:
                 print(f"Dockerfile not found at {dockerfile_path}")
                 sys.exit(1)
 
-            print(f"Using Dockerfile at: {dockerfile_path}")
+            self.debug(f"Using Dockerfile at: {dockerfile_path}")
 
             if not self.build_docker_image(docker_client, image_name, dockerfile_path, project_root):
                 print("Failed to build Docker image. Continuing without safe mode.")
@@ -108,7 +114,6 @@ class DockerController:
     @staticmethod
     def get_logs(container) -> str:
         # Stream the logs to the console
-        print("Running in Docker container...")
         logs = container.logs(stream=True)
         logs_output = ""
         for log in logs:
@@ -119,25 +124,32 @@ class DockerController:
 
     @staticmethod
     def run_container(docker_client: DockerClient, image_name: str, docker_args: list, project_root: str) -> Container:
+        # Create Docker-specific environment variables
+        docker_env = {
+            "RUNNING_IN_DOCKER": "1",
+            "PYTHONPATH": "/controller",
+            "COVERAGE_FILE": "/tmp/.coverage",  # Move coverage file to /tmp
+            "DB_PATH": "/tmp/testgen.db"        # Move DB to /tmp
+        }
+        
         return docker_client.containers.run(
                 image=image_name,
                 command=["poetry", "run", "python", "-m", "testgen.main"] + docker_args,
-                volumes={project_root: {"bind": "/controller", "mode": "rw"}},  # Mount current dir
-                environment={"RUNNING_IN_DOCKER": "1"},
+                volumes={project_root: {"bind": "/controller", "mode": "rw"}},
+                environment=docker_env,
                 detach=True,
                 remove=True,
                 stdout=True,
                 stderr=True
             )
 
-    @staticmethod
-    def build_docker_image(docker_client, image_name, dockerfile_path, project_root):
+    def build_docker_image(self, docker_client, image_name, dockerfile_path, project_root):
         try:
             print(f"Starting Docker build for image: {image_name}")
             dockerfile_rel_path = os.path.relpath(dockerfile_path, project_root)
-            print(f"Project root {project_root}")
-            print(f"Docker directory: {os.path.dirname(dockerfile_path)}")
-            print(f"Docker rel path: {dockerfile_rel_path}")
+            self.debug(f"Project root {project_root}")
+            self.debug(f"Docker directory: {os.path.dirname(dockerfile_path)}")
+            self.debug(f"Docker rel path: {dockerfile_rel_path}")
             build_progress = docker_client.api.build(
                 path=os.path.join(project_root, "testgen", "docker"),
                 dockerfile=os.path.join(project_root, "testgen", "docker", "Dockerfile"),
@@ -147,13 +159,13 @@ class DockerController:
             )
 
             for chunk in build_progress:
-                print(f"CHUNK: {chunk}")
+                self.debug(f"CHUNK: {chunk}")
                 if 'stream' in chunk:
                     for line in chunk['stream'].splitlines():
                         if line.strip():
                             print(f"Docker: {line.strip()}")
                 elif 'error' in chunk:
-                    print(f"Docker build error: {chunk['error']}")
+                    self.debug(f"Docker build error: {chunk['error']}")
                     return False
             print(f"Docker image built successfully: {image_name}")
             return True
@@ -166,4 +178,7 @@ class DockerController:
             print(f"Unexpected error during Docker build: {str(e)}")
             return False
 
-
+    def debug(self, message: str):
+        """Log debug message"""
+        if self.debug_mode:
+            self.logger.debug(message)

testgen/docker/Dockerfile

@@ -9,12 +9,12 @@ ENV POETRY_VIRTUALENVS_CREATE=false \
     PYTHONUNBUFFERED=1 \
     RUNNING_IN_DOCKER=true
 
-WORKDIR /app
+WORKDIR /controller
 
 # Copy poetry files
 COPY . .
 
-ENV PYTHONPATH=/app:/app/testgen
+ENV PYTHONPATH=/controller:/controller/testgen
 
 RUN poetry install --no-root
 

testgen/docker/pyproject.toml

@@ -1,10 +1,16 @@
 [tool.poetry]
-name = "testgen"
-version = "0.1.0"
+name = "testgenie-py"
+version = "0.1.6"
 description = ""
 authors = ["cjseitz <charlesjseitz@gmail.com>"]
 readme = "README.md"
 
+[[tool.poetry.packages]]
+include = "testgen" 
+
+[tool.poetry.scripts]
+testgenie = "testgen.main:main"
+
 [tool.poetry.dependencies]
 python = "^3.10"
 astor = "0.8.1"
@@ -16,6 +22,7 @@ typed-ast = "1.5.5"
 z3-solver = "4.13.3.0"
 staticfg = "^0.9.5"
 pytest = "^8.3.5"
+docker = "^7.1.0"
 
 [build-system]
 requires = ["poetry-core"]

testgen/reinforcement/environment.py

@@ -12,10 +12,11 @@ from testgen.models.test_case import TestCase
 
 
 class ReinforcementEnvironment:
-    def __init__(self, file_name, fut: ast.FunctionDef, module, initial_test_cases: List[TestCase], state: AbstractState):
+    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
         self.initial_test_cases = initial_test_cases
         self.test_cases = initial_test_cases.copy()
         self.state = state
@@ -32,13 +33,13 @@ class ReinforcementEnvironment:
 
         # Execute action
         if action == "add":
-            self.test_cases.append(randomizer.new_random_test_case(self.file_name, self.fut))
+            self.test_cases.append(randomizer.new_random_test_case(self.file_name, self.class_name, self.fut))
         elif action == "merge" and len(self.test_cases) > 1:
             self.test_cases.append(randomizer.combine_cases(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.fut, self.test_cases)
+            self.test_cases = randomizer.get_z3_test_cases(self.file_name, self.class_name, self.fut, self.test_cases)
         else:
             raise ValueError("Invalid action")
 
@@ -97,26 +98,52 @@ class ReinforcementEnvironment:
 
         test_cases = [tc for tc in self.test_cases if tc.func_name == self.fut.name]
 
+        executable_lines = set()
         if not test_cases:
             print("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.fut)
-            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.file_name, self.fut.name, temp_case.inputs)
+            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.fut.name, test_cases[0].inputs)
+            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(analysis[1])
+        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)
-        
+        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):
@@ -153,7 +180,12 @@ class ReinforcementEnvironment:
         for test_case in self.test_cases:
             try:
                 module = testgen.util.file_utils.load_module(self.file_name)
-                func = getattr(module, self.fut.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(*test_case.inputs)
             except Exception as e:
                 import traceback

testgen/reinforcement/statement_coverage_state.py

@@ -12,7 +12,7 @@ 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.fut.name, test_case.inputs)
+            analysis = testgen.util.coverage_utils.get_coverage_analysis(self.environment.file_name, self.environment.class_name, self.environment.fut.name, test_case.inputs)
             covered = testgen.util.coverage_utils.get_list_of_covered_statements(analysis)
             all_covered_statements.update(covered)
 

testgen/service/analysis_service.py

@@ -43,10 +43,15 @@ class AnalysisService:
     
     def create_analysis_context(self, filepath: str) -> AnalysisContext:
         """Create an analysis context for the given file."""
+        print(f"Creating analysis context for {filepath}")
         filename = file_utils.get_filename(filepath)
+        print(f"Filename: {filename}")
         module = file_utils.load_module(filepath)
+        print(f"Module: {module}")
         class_name = self.get_class_name(module)
+        print(f"Class name: {class_name}")
         function_data = self.get_function_data(filename, module, class_name)
+        print(f"Function data: {function_data}")
         return AnalysisContext(filepath, filename, class_name, module, function_data)
     
     def get_function_data(self, filename: str, module: ModuleType, class_name: str | None) -> List[FunctionMetadata]:
@@ -77,11 +82,12 @@ class AnalysisService:
 
         return function_metadata_list
 
-    def do_reinforcement_learning(self, filepath: str, mode: str = None) -> List[TestCase]:
+    def do_reinforcement_learning(self, filepath: str, class_name: str | None, 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)
+        self.class_name = class_name
         time_limit: int = 30
         all_test_cases: List[TestCase] = []
 
@@ -94,7 +100,7 @@ class AnalysisService:
             best_coverage: float = 0.0
 
             # Create environment and agent once per function
-            environment = ReinforcementEnvironment(filepath, function, module, function_test_cases, state=StatementCoverageState(None))
+            environment = ReinforcementEnvironment(filepath, function, module, self.class_name, function_test_cases, state=StatementCoverageState(None))
             environment.state = StatementCoverageState(environment)
 
             # Create agent with existing Q-table

testgen/service/cfg_service.py

@@ -44,7 +44,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, func.func_def, test_cases)
+            self.visualizer.get_covered_lines(file_path, analysis_context.class_name, func.func_def, test_cases)
 
             base_filename = f"{filename}_{func.function_name}_coverage"
             output_filepath = self.get_versioned_filename(visualization_dir, base_filename)

testgen/service/generator_service.py

@@ -1,4 +1,5 @@
 import os
+import inspect
 from types import ModuleType
 from typing import List
 
@@ -43,15 +44,22 @@ class GeneratorService:
 
     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)
         
+        # Determine the actual class name used in the module
+        actual_class_name = class_name
+        if 'generated_' in self.filepath and class_name:
+            # For generated classes, find the actual class name in the module
+            for name, obj in inspect.getmembers(module):
+                if inspect.isclass(obj):
+                    actual_class_name = name
+                    break
+        
         context = GeneratorContext(
             filepath=self.filepath,
             filename=filename,
-            class_name=class_name,
+            class_name=actual_class_name,  # Use the actual class name
             module=module,
             output_path=output_path,
             test_cases=test_cases