testgenie-py 0.1.0__py3-none-any.whl

testgen/analyzer/test_case_analyzer.py

@@ -0,0 +1,46 @@
+import ast
+from abc import ABC, abstractmethod
+from typing import List, Dict
+
+from testgen.models.test_case import TestCase
+from testgen.models.analysis_context import AnalysisContext
+from testgen.models.function_metadata import FunctionMetadata
+
+class TestCaseAnalyzerStrategy(ABC):
+    def __init__(self, analysis_context: AnalysisContext = None):
+        self._analysis_context = analysis_context
+
+    @abstractmethod
+    def collect_test_cases(self, function_metadata: FunctionMetadata) -> List[TestCase]:
+        pass
+
+    def get_function_metadata(self, func_name: str) -> FunctionMetadata | None:
+        if self._analysis_context and self._analysis_context.function_data:
+            for func_metadata in self._analysis_context.function_data:
+                if func_metadata.function_name == func_name:
+                    return func_metadata
+        return None
+
+    def get_function_metadata_by_node(self, func_node: ast.FunctionDef) -> FunctionMetadata:
+        return self.get_function_metadata(func_node.name)
+
+    def get_param_types(self, func_node: ast.FunctionDef) -> Dict[str, str] | None:
+        func_metadata = self.get_function_metadata_by_node(func_node)
+        if func_metadata:
+            return func_metadata.params
+
+        return None
+
+    @property
+    def analysis_context(self) -> AnalysisContext:
+        return self._analysis_context
+        
+    @analysis_context.setter
+    def analysis_context(self, context: AnalysisContext):
+        self._analysis_context = context
+        
+
+    
+
+    
+

testgen/analyzer/test_case_analyzer_context.py

@@ -0,0 +1,58 @@
+import time
+from typing import List
+
+from testgen.analyzer.test_case_analyzer import TestCaseAnalyzerStrategy
+from testgen.models.test_case import TestCase
+from testgen.models.analysis_context import AnalysisContext
+
+
+class TestCaseAnalyzerContext:
+    def __init__(self, analysis_context: AnalysisContext, test_case_analyzer: TestCaseAnalyzerStrategy):
+        self._test_case_analyzer = test_case_analyzer
+        self._analysis_context = analysis_context
+        self._test_cases = []
+        
+    # TODO: GET RID OF THIS STUPID METHOD IT IS POINTLESS
+    # JUST CALL INSIDE ANALYZER_SERVICE
+    def do_logic(self) -> List[TestCase]:
+        """Run the analysis process"""
+        self.do_strategy(20)
+
+    def do_strategy(self, time_limit=None) -> List[TestCase]:
+        """Execute the analysis strategy for all functions with an optional time limit"""
+        start_time = time.time()
+        for func_metadata in self._analysis_context.function_data:
+            if time_limit and (time.time() - start_time) >= time_limit:
+                break
+            test_cases = self._test_case_analyzer.collect_test_cases(func_metadata)
+            for test_case in test_cases:
+                print(f"Test Case: {test_case.func_name}, {test_case.inputs}, {test_case.expected}")
+            self._test_cases.extend(test_cases)
+
+    @property
+    def strategy(self) -> TestCaseAnalyzerStrategy:
+        return self._test_case_analyzer
+
+    @property
+    def test_cases(self):
+        return self._test_cases
+
+    @strategy.setter
+    def strategy(self, test_case_analyzer: TestCaseAnalyzerStrategy) -> None:
+        self._test_case_analyzer = test_case_analyzer
+        if self._analysis_context:
+            self._test_case_analyzer.analysis_context = self._analysis_context
+
+    @property
+    def analysis_context(self) -> AnalysisContext:
+        return self._analysis_context
+    
+    @analysis_context.setter
+    def analysis_context(self, context: AnalysisContext) -> None:
+        self._analysis_context = context
+        if self._test_case_analyzer:
+            self._test_case_analyzer.analysis_context = context
+
+    
+
+

testgen/controller/cli_controller.py

@@ -0,0 +1,194 @@
+import argparse
+import inspect
+import os
+import sys
+
+import docker
+from docker import DockerClient
+from docker import errors
+
+from testgen.controller.docker_controller import DockerController
+from testgen.service.service import Service
+from testgen.presentation.cli_view import CLIView
+from testgen.sqlite.db_service import DBService
+
+AST_STRAT = 1
+FUZZ_STRAT = 2
+RANDOM_STRAT = 3
+REINFORCE_STRAT = 4
+
+UNITTEST_FORMAT = 1
+PYTEST_FORMAT = 2
+DOCTEST_FORMAT = 3
+
+class CLIController:
+    #TODO: Possibly create a view 'interface' and use dependency injection to extend other views
+    def __init__(self, service: Service, view: CLIView):
+        self.service = service
+        self.view = view
+
+    def run(self):
+        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.")
+        parser.add_argument(
+            "--generate-only", "-g",
+            action="store_true",
+            help="Generate branched code but skip running unit tests and coverage."
+        )
+        parser.add_argument(
+            "--test-mode",
+            choices=["ast", "random", "fuzz", "reinforce"],
+            default="ast",
+            help="Set the test generation analysis technique"
+        )
+        parser.add_argument(
+            "--reinforce-mode",
+            choices=["train", "collect"],
+            default="train",
+            help="Set mode for reinforcement learning"
+        )
+        parser.add_argument(
+            "--test-format",
+            choices=["unittest", "pytest", "doctest"],
+            default="unittest",
+            help="Set the test generation format"
+        )
+        parser.add_argument(
+            "--safe",
+            action="store_true",
+            help="Run test generation from within a docker container."
+        )
+        parser.add_argument(
+            "--db",
+            type=str,
+            default="testgen.db",
+            help="Path to SQLite database file (default: testgen.db)"
+        )
+        parser.add_argument(
+            "--select-all",
+            action="store_true",
+            help="Select all from sqlite db"
+        )
+        parser.add_argument(
+            "--visualize",
+            action="store_true",
+            help = "Visualize the tests with graphviz"
+        )
+
+        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)
+        else:
+            self.view.display_message("Running in local mode...")
+            self.execute_generation(args)
+
+    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)
+            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
+
+    def docker_available(self) -> DockerClient | None:
+        try:
+            client = docker.from_env()
+            client.ping()
+            self.view.display_message("Docker daemon is running and connected.")
+            return client
+        except docker.errors.DockerException as err:
+            print(f"Docker is not available: {err}")
+            print(f"Make sure the Docker daemon is running, and try again.")
+            choice = input("Continue without Docker (y/n)?")
+            if choice.lower() == 'y':
+                return None
+            else:
+                sys.exit(1)

testgen/controller/docker_controller.py

@@ -0,0 +1,169 @@
+import os
+import sys
+from argparse import Namespace
+import docker
+from docker import DockerClient, client
+from docker import errors
+from docker.models.containers import Container
+
+from testgen.service.service import Service
+
+AST_STRAT = 1
+FUZZ_STRAT = 2
+RANDOM_STRAT = 3
+
+UNITTEST_FORMAT = 1
+PYTEST_FORMAT = 2
+DOCTEST_FORMAT = 3
+
+class DockerController:
+    def __init__(self):
+        self.service = Service()
+        self.args = None
+
+    def run_in_docker(self, project_root: str, docker_client: DockerClient, args: Namespace) -> bool:
+        self.args = args
+        os.environ["RUNNING_IN_DOCKER"] = "1"
+
+        # Check if Docker image exists, build it if not
+        image_name = "testgen-runner"
+        # 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...")
+            return False
+
+        docker_args = [os.path.basename(args.file_path)] + [arg for arg in sys.argv[2:] if arg != "--safe"]
+
+        # Run the container with the same arguments
+        try:
+            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)
+
+            try:
+                # Create the target directory if it doesn't exist
+                if args.output is None:
+                    target_path = os.path.join(os.getcwd(), 'tests')
+                else:
+                    target_path = args.output
+                os.makedirs(target_path, exist_ok=True)
+
+                print(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.service.set_file_path(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)
+
+                test_file = self.service.generate_test_file(test_cases, target_path)
+                print(f"Unit tests saved to: {test_file}")
+
+                if not args.generate_only:
+                    print("Running coverage...")
+                    import traceback
+                    print(traceback.format_exc())
+                    self.service.run_coverage(test_file)
+
+            except Exception as e:
+                print(f"Error running container: {e}")
+                sys.exit(1)
+
+        except Exception as e:
+            print(f"Error running container: {e}")
+            sys.exit(1)
+
+    def get_image(self, docker_client: DockerClient, image_name: str, project_root: str):
+        try:
+            docker_client.images.get(image_name)
+            print(f"Using existing {image_name} Docker image")
+        except docker.errors.ImageNotFound:
+            print(f"Building {image_name} Docker image...")
+
+            # Look for Dockerfile in the project root
+            docker_dir = os.path.join(project_root, "testgen", "docker")
+            dockerfile_path = os.path.join(docker_dir, "Dockerfile")
+            if not os.path.exists(dockerfile_path):
+                print(f"Dockerfile not found at {dockerfile_path}")
+                sys.exit(1)
+
+            print(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.")
+                self.args.safe = False
+
+    @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:
+            log_line = log.decode()
+            logs_output += log_line
+            print(log_line, end="")
+        return logs_output
+
+    @staticmethod
+    def run_container(docker_client: DockerClient, image_name: str, docker_args: list, project_root: str) -> Container:
+        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"},
+                detach=True,
+                remove=True,
+                stdout=True,
+                stderr=True
+            )
+
+    @staticmethod
+    def build_docker_image(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}")
+            build_progress = docker_client.api.build(
+                path=os.path.join(project_root, "testgen", "docker"),
+                dockerfile=os.path.join(project_root, "testgen", "docker", "Dockerfile"),
+                tag=image_name,
+                rm=True,
+                decode=True
+            )
+
+            for chunk in build_progress:
+                print(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']}")
+                    return False
+            print(f"Docker image built successfully: {image_name}")
+            return True
+
+        except docker.errors.BuildError as e:
+            print(f"Docker build error: {e}")
+        except docker.errors.APIError as e:
+            print(f"Docker API error: {e}")
+        except Exception as e:
+            print(f"Unexpected error during Docker build: {str(e)}")
+            return False
+
+

testgen/docker/Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.10
+
+RUN apt-get update && apt-get install -y curl build-essential
+
+RUN curl -sSL https://install.python-poetry.org | python3 - \
+    && ln -s /root/.local/bin/poetry /usr/local/bin/poetry
+
+ENV POETRY_VIRTUALENVS_CREATE=false \
+    PYTHONUNBUFFERED=1 \
+    RUNNING_IN_DOCKER=true
+
+WORKDIR /app
+
+# Copy poetry files
+COPY . .
+
+ENV PYTHONPATH=/app:/app/testgen
+
+RUN poetry install --no-root
+
+# Entrypoint (Can be overridden by Docker SDK in Python)
+CMD ["poetry", "run", "python", "-m", "testgen.main"]