devdox-ai-locust 0.1.1__py3-none-any.whl

devdox_ai_locust/locust_generator.py

@@ -0,0 +1,732 @@
+"""
+Locust Test Generator
+
+Generates Locust performance test files from parsed OpenAPI endpoints.
+"""
+
+import json
+import re
+import secrets
+from typing import List, Tuple, Dict, Any, Optional
+import black
+from jinja2 import Environment, FileSystemLoader
+from pathlib import Path
+import logging
+from dataclasses import dataclass
+from datetime import datetime
+
+
+from devdox_ai_locust.utils.open_ai_parser import Endpoint, Parameter
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class TestDataConfig:
+    """Configuration for test data generation"""
+
+    string_length: int = 10
+    integer_min: int = 1
+    integer_max: int = 1000
+    array_size: int = 3
+    use_realistic_data: bool = True
+
+
+class LocustTestGenerator:
+    """Generates Locust performance test files from OpenAPI endpoints"""
+
+    def __init__(
+        self,
+        test_config: Optional[TestDataConfig] = None,
+        jinja_env: Optional[Environment] = None,
+    ):
+        self.test_config = test_config or TestDataConfig()
+        self.generated_files: Dict[str, str] = {}
+        self.auth_token = None
+        self.user_data: Dict[str, Any] = {}
+        self.request_count = 0
+
+        self.template_dir = self._find_project_root() / "templates"
+
+        self.template_dir.mkdir(parents=True, exist_ok=True)
+
+        self.jinja_env = jinja_env or self._setup_templates()
+
+    def _find_project_root(self) -> Path:
+        """Find the project root by looking for setup.py, pyproject.toml, or .git"""
+        current_path = Path(__file__).parent
+
+        return current_path
+
+    def _setup_templates(self) -> Environment:
+        """Initialize Jinja2 environment and create default templates"""
+        return Environment(
+            loader=FileSystemLoader(self.template_dir),
+            trim_blocks=True,
+            lstrip_blocks=True,
+            keep_trailing_newline=True,
+            autoescape=False,
+        )
+
+    def fix_indent(self, base_files: Dict[str, str]) -> Dict[str, str]:
+        """Fix indentation for generated files"""
+        try:
+            mode = black.Mode()
+            updated_data = {}
+
+            for key, value in base_files.items():
+                try:
+                    formatted_code = black.format_str(value, mode=mode)
+
+                    updated_data[key] = formatted_code
+                except black.InvalidInput:
+                    # Not valid Python code, keep original
+                    logger.debug(f"Skipping formatting for {key}: not valid Python")
+                    updated_data[key] = value
+                except Exception as format_error:
+                    # Other Black formatting errors, keep original and log
+                    logger.warning(f"Failed to format {key}: {format_error}")
+                    updated_data[key] = value
+
+            return updated_data
+
+        except Exception as e:
+            logger.error(f"exception occurred: {e}")
+            return base_files
+
+    def generate_from_endpoints(
+        self,
+        endpoints: List[Endpoint],
+        api_info: Dict[str, Any],
+        include_auth: bool = True,
+        target_host: Optional[str] = None,
+    ) -> Tuple[Dict[str, str], List[Dict[str, Any]], Dict[str, List[Endpoint]]]:
+        """
+        Generate complete Locust test suite from parsed endpoints
+
+        Args:
+            endpoints: List of parsed Endpoint objects
+            api_info: API information dictionary
+            output_dir: Output directory for generated files
+
+        Returns:
+            Dictionary of filename -> file content
+        """
+        try:
+            grouped_enpoint = self._group_endpoints_by_tag(endpoints, include_auth)
+
+            workflows_files = self.generate_workflows(grouped_enpoint, api_info)
+
+            self.generated_files = {
+                "locustfile.py": self._generate_main_locustfile(
+                    endpoints, api_info, list(grouped_enpoint.keys())
+                ),
+                "test_data.py": self._generate_test_data_file(),
+                "config.py": self._generate_config_file(api_info),
+                "utils.py": self._generate_utils_file(),
+                "custom_flows.py": self._generate_custom_flows_file(),
+                "requirements.txt": self._generate_requirements_file(),
+                "README.md": self._generate_readme_file(api_info),
+                ".env.example": self._generate_env_example(api_info, target_host),
+            }
+
+            return self.generated_files, workflows_files, grouped_enpoint
+        except Exception as e:
+            logger.error(f"Failed to generate test suite: {e}")
+            empty_files: Dict[str, str] = {}
+            empty_workflows: List[Dict[str, Any]] = []
+            empty_grouped: Dict[str, List[Endpoint]] = {}
+            return empty_files, empty_workflows, empty_grouped
+
+    def generate_workflows(
+        self, endpoints: Dict[str, List[Any]], api_info: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """
+        Generate the workflows Locust test files with proper structure and no duplicates
+        """
+        try:
+            workflows: List = []
+
+            for group, group_endpoints in endpoints.items():
+                task_methods: List[str] = []
+
+                for endpoint in group_endpoints:
+                    try:
+                        task_method = self._generate_task_method(endpoint)
+                        if task_method:
+                            task_methods.append(task_method)
+                    except Exception as e:
+                        logger.warning(
+                            f"⚠️ Failed to generate task method for {getattr(endpoint, 'path', '?')}: {e}"
+                        )
+                        continue
+                if not task_methods:
+                    logger.warning(f"No task methods generated from group {group}")
+                    task_methods.append(self._generate_default_task_method())
+                indented_task_methods = self._indent_methods(
+                    task_methods, indent_level=1
+                )
+                file_content = self._build_endpoint_template(
+                    api_info, indented_task_methods, group
+                )
+                file_name = f"{group}_workflow.py".replace("-", "_")
+
+                workflows.append({file_name: file_content})
+
+            workflows.append(
+                {"base_workflow.py": self.generate_base_common_file(api_info)}
+            )
+
+            return workflows
+
+        except Exception as e:
+            logger.error(f"❌ Failed to generate test suite: {e}")
+
+            return []
+
+    def _build_endpoint_template(
+        self, api_info: Dict[str, Any], task_methods_content: str, group: str
+    ) -> str:
+        template = self.jinja_env.get_template("endpoint_template.py.j2")
+        return template.render(
+            api_info=api_info, group=group, task_methods_content=task_methods_content
+        )
+
+    def _generate_main_locustfile(
+        self, endpoints: List[Any], api_info: Dict[str, Any], groups: List[str]
+    ) -> str:
+        """
+        Generate the main Locust test file with proper structure and no duplicates
+
+        Args:
+            endpoints: List of parsed Endpoint objects
+            api_info: API information dictionary
+
+        Returns:
+            Complete locustfile.py content as string
+        """
+        try:
+            # Generate task methods for each endpoint
+            task_methods = []
+            for endpoint in endpoints:
+                try:
+                    task_method = self._generate_task_method(endpoint)
+                    if task_method:
+                        task_methods.append(task_method)
+                except Exception as e:
+                    logger.warning(
+                        f"Failed to generate task method for  {endpoint.path}: {e}"
+                    )
+                    continue
+
+            if not task_methods:
+                logger.warning("No task methods generated from endpoints")
+                # Generate a default task method
+                task_methods.append(self._generate_default_task_method())
+
+            # Properly indent task methods for class inclusion
+            indented_task_methods = self._indent_methods(task_methods, indent_level=1)
+            indented_task_methods = ""
+            # Generate the complete file content
+            return self._build_locustfile_template(
+                api_info=api_info,
+                task_methods_content=indented_task_methods,
+                groups=groups,
+            )
+
+        except Exception as e:
+            logger.error(f"Failed to generate test suite: {e}")
+            # Return fallback files
+            return self._generate_fallback_locustfile(api_info)
+
+    def _indent_methods(self, task_methods: List[str], indent_level: int = 1) -> str:
+        """Properly indent task methods for class inclusion"""
+        indented_methods = []
+        for method in task_methods:
+            lines = method.split("\n")
+            indented_lines = []
+            first_nonempty = True
+
+            for line in lines:
+                if line.strip():
+                    stripped_line = line.lstrip()
+                    if first_nonempty:
+                        # First non-empty line → indent once (for @task or def)
+                        indented_lines.append("    " * indent_level + stripped_line)
+                        first_nonempty = False
+                    else:
+                        # Method body → indent deeper
+                        indented_lines.append(
+                            "    " * (indent_level + 1) + stripped_line
+                        )
+                else:
+                    indented_lines.append("")
+            indented_methods.append("\n".join(indented_lines))
+
+        return "\n\n".join(indented_methods)
+
+    def _generate_fallback_locustfile(self, api_info: Dict[str, Any]) -> str:
+        """Generate a basic fallback locustfile when main generation fails"""
+
+        template = self.jinja_env.get_template("fallback_locust.py.j2")
+        return template.render(api_info=api_info)
+
+    from typing import Dict, Any
+
+    def _build_locustfile_template(
+        self, api_info: Dict[str, Any], task_methods_content: str, groups: List[str]
+    ) -> str:
+        import_group_tasks = ""
+        tasks = []
+        for group in groups:
+            file_name = group.lower().replace("-", "_")
+            class_name = group.replace("-", "")
+            import_group_tasks += f"""from workflows.{file_name}_workflow import {class_name}TaskMethods\n"""
+            tasks.append(f"{class_name}TaskMethods")
+        tasks_str = "[" + ",".join(tasks) + "]"
+        template = self.jinja_env.get_template("locust.py.j2")
+
+        # Prepare template context
+        context = {
+            "import_group_tasks": import_group_tasks,
+            "task_methods_content": task_methods_content,
+            "tasks_str": tasks_str,
+            "api_info": api_info,
+            "generated_task_classes": self._generate_user_classes(),
+            "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+        }
+
+        # Render the template
+        content = template.render(**context)
+
+        logger.info("README generated successfully using template")
+        return content
+
+    def _generate_default_task_method(self) -> str:
+        """Generate a default task method when no endpoints are available"""
+        return '''@task(1)
+    def default_health_check(self):
+        """Default health check task"""
+        try:
+            response_data = self.make_request(
+                method="get",
+                path="/health"
+            )
+
+            if response_data:
+                self._store_response_data("health_check", response_data)
+
+        except Exception as e:
+            logger.error(f"Health check task failed: {e}")
+    '''
+
+    def _generate_task_method(self, endpoint: Any) -> str:
+        """Generate a Locust task method for a single endpoint with improved structure"""
+        try:
+            method_name = self._generate_method_name(endpoint)
+            path_with_params = self._generate_path_with_params(endpoint)
+            weight = self._get_task_weight(getattr(endpoint, "method", "GET"))
+
+            # Build the task method with proper error handling
+            task_method = f'''@task({weight})
+    def {method_name}(self):
+        """
+        {getattr(endpoint, "summary", f"{getattr(endpoint, 'method', 'GET')} {getattr(endpoint, 'path', '')}")}
+        {getattr(endpoint, "description", "")}
+        """
+        try:
+            # Generate path parameters
+            {self._generate_path_params_code(endpoint)}
+
+            # Generate query parameters
+            {self._generate_query_params_code(endpoint)}
+
+            # Generate request body
+            {self._generate_request_body_code(endpoint)}
+
+            # Make the request
+            response_data = self.make_request(
+                method="{getattr(endpoint, "method", "GET").lower()}",
+                path=f"{path_with_params}",
+                {self._generate_request_kwargs(endpoint)}
+            )
+
+            if response_data:
+                # Store response data for dependent requests
+                self._store_response_data("{method_name}", response_data)
+
+        except Exception as e:
+            logger.error(f"Task {method_name} failed: {{e}}")
+    '''
+            return task_method
+
+        except Exception as e:
+            logger.error(f"Failed to generate task method for 535 endpoint: {e}")
+            return ""
+
+    def _generate_method_name(self, endpoint: Endpoint) -> str:
+        """Generate a valid Python method name from endpoint"""
+        if endpoint.operation_id:
+            # Use operation ID if available
+            name = endpoint.operation_id
+        else:
+            # Generate name from method and path
+            path_parts = [
+                part
+                for part in endpoint.path.split("/")
+                if part and not part.startswith("{")
+            ]
+            name = f"{endpoint.method.lower()}_{'_'.join(path_parts)}"
+
+        # Clean up the name
+        name = re.sub(r"[^\w]", "_", name)
+        name = re.sub(r"_+", "_", name)
+        name = name.strip("_")
+
+        return name if name else f"{endpoint.method.lower()}_endpoint"
+
+    def _generate_path_with_params(self, endpoint: Endpoint) -> str:
+        """Generate path with parameter placeholders"""
+        path = endpoint.path
+
+        # Replace path parameters with f-string format
+        for param in endpoint.parameters:
+            if param.location.value == "path":
+                path = path.replace(f"{{{param.name}}}", f"{{{param.name}}}")
+
+        return path
+
+    def _generate_path_params_code(self, endpoint: Endpoint) -> str:
+        """Generate code for path parameters"""
+        path_params = [p for p in endpoint.parameters if p.location.value == "path"]
+
+        if not path_params:
+            return "# No path parameters"
+
+        code_lines = []
+        for param in path_params:
+            if param.type.startswith("integer"):
+                code_lines.append(f"{param.name} = data_generator.generate_integer()")
+            elif param.type == "string":
+                if "id" in param.name.lower():
+                    code_lines.append(f"{param.name} = data_generator.generate_id()")
+                else:
+                    code_lines.append(
+                        f"{param.name} = data_generator.generate_string()"
+                    )
+
+            else:
+                code_lines.append(
+                    f'{param.name} = data_generator.generate_value("{param.type}")'
+                )
+
+        return "\n".join(code_lines)
+
+    def _generate_query_params_code(self, endpoint: Endpoint) -> str:
+        """Generate code for query parameters"""
+        query_params = [p for p in endpoint.parameters if p.location.value == "query"]
+
+        if not query_params:
+            return "params = {}"
+
+        param_lines = []
+        for param in query_params:
+            if self._should_skip_optional_param(param):
+                continue
+
+            param_line = self._generate_param_line(param)
+            param_lines.append(param_line)
+
+        return self._format_params_dict(param_lines)
+
+    def _should_skip_optional_param(self, param: Parameter) -> bool:
+        """Randomly skip optional parameters 30% of the time"""
+        return not param.required and secrets.randbelow(100) > 70
+
+    def _generate_param_line(self, param: Parameter) -> str:
+        """Generate a single parameter line based on its type"""
+        param_generators = {
+            "integer": self._generate_integer_param,
+            "string": self._generate_string_param,
+            "boolean": self._generate_boolean_param,
+        }
+
+        # Handle integer types that might have prefixes like "integer64"
+        param_type = "integer" if param.type.startswith("integer") else param.type
+
+        generator = param_generators.get(param_type, self._generate_generic_param)
+        return generator(param)
+
+    def _generate_integer_param(self, param: Parameter) -> str:
+        """Generate integer parameter line"""
+        default = param.default if param.default is not None else "None"
+        return f'"{param.name}": data_generator.generate_integer(default={default}),'
+
+    def _generate_string_param(self, param: Parameter) -> str:
+        """Generate string parameter line"""
+        default = f'"{param.default}"' if param.default else "None"
+        return f'"{param.name}": data_generator.generate_string(default={default}),'
+
+    def _generate_boolean_param(self, param: Parameter) -> str:
+        """Generate boolean parameter line"""
+        return f'"{param.name}": data_generator.generate_boolean(),'
+
+    def _generate_generic_param(self, param: Parameter) -> str:
+        """Generate generic parameter line for unknown types"""
+        return f'"{param.name}": data_generator.generate_value("{param.type}"),'
+
+    def _format_params_dict(self, param_lines: list[str]) -> str:
+        """Format parameter lines into a dictionary structure"""
+        if not param_lines:
+            return "params = {}"
+
+        lines = ["params = {"] + [f"    {line}" for line in param_lines] + ["}"]
+        return "\n".join(lines)
+
+    def _generate_request_body_code(self, endpoint: Endpoint) -> str:
+        """Generate code for request body"""
+        if not endpoint.request_body:
+            return "json_data = None"
+
+        if endpoint.request_body.content_type == "application/json":
+            return f"""json_data = data_generator.generate_json_data(
+                schema={json.dumps(endpoint.request_body.schema, indent=16)}
+            )"""
+        elif endpoint.request_body.content_type == "application/x-www-form-urlencoded":
+            return "data = data_generator.generate_form_data()"
+        else:
+            return "json_data = {}"
+
+    def _generate_request_kwargs(self, endpoint: Endpoint) -> str:
+        """Generate kwargs for the request method"""
+        kwargs = []
+
+        # Add query parameters
+        query_params = [p for p in endpoint.parameters if p.location.value == "query"]
+        if query_params:
+            kwargs.append("params=params")
+
+        # Add request body
+        if endpoint.request_body:
+            if endpoint.request_body.content_type == "application/json":
+                kwargs.append("json=json_data")
+            elif (
+                endpoint.request_body.content_type
+                == "application/x-www-form-urlencoded"
+            ):
+                kwargs.append("data=data")
+
+        # Add headers if needed
+        header_params = [p for p in endpoint.parameters if p.location.value == "header"]
+        if header_params:
+            kwargs.append("headers=headers")
+
+        return ",\n                ".join(kwargs)
+
+    def _get_task_weight(self, method: str) -> int:
+        """Get task weight based on HTTP method"""
+        weights = {
+            "GET": 5,  # Most frequent
+            "POST": 2,  # Common
+            "PUT": 1,  # Less frequent
+            "PATCH": 1,  # Less frequent
+            "DELETE": 1,  # Least frequent
+            "HEAD": 3,  # Moderate
+            "OPTIONS": 1,  # Rare
+        }
+        return weights.get(method.upper(), 1)
+
+    def _group_endpoints_by_tag(
+        self,
+        endpoints: List[Endpoint],
+        include_auth_endpoints: bool = True,
+    ) -> Dict[str, List[Endpoint]]:
+        """Group endpoints by their tags"""
+        grouped: Dict[str, List[Endpoint]] = {}
+        # Define authentication-related keywords to check for in paths
+        auth_keywords = [
+            "login",
+            "signin",
+            "sign-in",
+            "sign_in",
+            "logout",
+            "signout",
+            "sign-out",
+            "sign_out",
+            "auth",
+            "authenticate",
+            "authorization",
+            "token",
+            "refresh",
+            "verify",
+            "password",
+            "reset",
+            "forgot",
+            "register",
+            "signup",
+            "sign-up",
+            "sign_up",
+            "session",
+            "oauth",
+            "sso",
+        ]
+
+        def is_auth_endpoint(endpoint_path: str) -> bool:
+            """Check if endpoint path contains authentication-related keywords"""
+            path_lower = endpoint_path.lower()
+            return any(keyword in path_lower for keyword in auth_keywords)
+
+        for endpoint in endpoints:
+            tags = endpoint.tags if endpoint.tags else ["default"]
+            if is_auth_endpoint(endpoint.path) and include_auth_endpoints:
+                # Add to authentication group regardless of tags
+                if "Authentication" not in grouped:
+                    grouped["Authentication"] = []
+                grouped["Authentication"].append(endpoint)
+
+            for tag in tags:
+                if tag not in grouped:
+                    grouped[tag] = []
+                grouped[tag].append(endpoint)
+
+        return grouped
+
+    def _generate_all_task_methods_string(self, endpoints: List[Endpoint]) -> str:
+        """Generate all task methods as a properly indented string"""
+        methods = []
+        for endpoint in endpoints:
+            method_code = self._generate_task_method(endpoint)
+            methods.append(method_code)
+
+        return "\n".join(methods)
+
+    def _generate_user_classes(self) -> str:
+        """
+        **FIXED: Generate user classes with proper structure**
+        """
+
+        return '''
+    class LightUser(BaseAPIUser, BaseTaskMethods):
+        """Light user with occasional API usage patterns"""
+        weight = 3
+        wait_time = between(3, 8)  # Longer wait times
+
+        def on_start(self):
+            super().on_start()
+            self.user_type = "light"
+   
+
+    class RegularUser( BaseAPIUser, BaseTaskMethods):
+        """Regular user with normal API usage patterns"""
+        weight = 4
+        wait_time = between(1, 4)  # Moderate wait times
+        
+        def on_start(self):
+            super().on_start()
+            self.user_type = "regular"
+        
+
+
+    class PowerUser( BaseAPIUser, BaseTaskMethods):
+        """Power user with heavy API usage patterns"""
+        weight = 3
+        wait_time = between(0.5, 2)  # Shorter wait times
+
+        def on_start(self):
+            super().on_start()
+            self.user_type = "power"
+            
+
+    '''
+
+    def _generate_test_data_file(self) -> str:
+        """Generate test_data.py file content"""
+        template = self.jinja_env.get_template("test_data.py.j2")
+        return template.render()
+
+    def _generate_config_file(self, api_info: Dict[str, Any]) -> str:
+        """Generate config.py file content"""
+        template = self.jinja_env.get_template("config.py.j2")
+        return template.render(api_info=api_info)
+
+    def _generate_utils_file(self) -> str:
+        """Generate utils.py file content"""
+        template = self.jinja_env.get_template("utils.py.j2")
+        return template.render()
+
+    def _generate_custom_flows_file(self) -> str:
+        """Generate custom_flows.py file content"""
+        template = self.jinja_env.get_template("custom_flows.py.j2")
+        return template.render()
+
+    def _generate_requirements_file(self) -> str:
+        """Generate requirements.txt file content"""
+        template = self.jinja_env.get_template("requirement.txt.j2")
+        content = template.render()
+        return content
+
+    def _generate_env_example(
+        self, api_info: Dict[str, Any], target_host: Optional[str] = None
+    ) -> str:
+        """Generate .env.example file content"""
+        try:
+            template = self.jinja_env.get_template("env.example.j2")
+            if target_host:
+                base_url = target_host
+                locust_host = target_host
+            else:
+                base_url = api_info.get("base_url", "http://localhost:8000")
+                locust_host = api_info.get("base_url", "http://localhost:8000")
+            # Prepare environment variables context
+            environment_vars = {
+                "API_BASE_URL": base_url,
+                "API_VERSION": api_info.get("version", "v1"),
+                "API_TITLE": api_info.get("title", "Your API Name"),
+                "LOCUST_USERS": "50",
+                "LOCUST_SPAWN_RATE": "5",
+                "LOCUST_RUN_TIME": "10m",
+                "LOCUST_HOST": locust_host,
+                "USE_REALISTIC_DATA": "true",
+                "DATA_SEED": "42",
+                "REQUEST_TIMEOUT": "30",
+                "MAX_RETRIES": "3",
+            }
+
+            context = {
+                "environment_vars": environment_vars,
+                "api_info": api_info,
+                "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            }
+
+            content = template.render(**context)
+            logger.info("✅ .env.example generated successfully using template")
+            return content
+
+        except Exception as e:
+            logger.error(f"❌ Failed to generate .env.example from template: {e}")
+            return ""
+
+    def _generate_readme_file(self, api_info: Dict[str, Any]) -> str:
+        try:
+            # Get the template
+            template = self.jinja_env.get_template("readme.md.j2")
+
+            # Prepare template context
+            context = {
+                "api_info": api_info,
+                "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+            }
+
+            # Render the template
+            content = template.render(**context)
+
+            logger.info("README generated successfully using template")
+            return content
+
+        except Exception as e:
+            logger.error(f"Error generating README: {e}")
+            return ""
+
+    def generate_base_common_file(self, api_info: Dict[str, Any]) -> str:
+        template = self.jinja_env.get_template("base_workflow.py.j2")
+        return template.render(api_info=api_info)