From 87e2e9fa79c6e92ceb023ad8cbc1497e04aa526a Mon Sep 17 00:00:00 2001
From: Subham Agrawal <subhamagrawal@Subhams-MacBook-Pro.local>
Date: Fri, 8 Aug 2025 23:43:22 +0530
Subject: [PATCH 1/3] feat: Add plugin system for custom request functions

- Introduced a new plugin system allowing users to define custom request functions using @request and @on_call_completion decorators.
- Implemented a PluginLoader to automatically discover and load decorated functions from Python files.
- Updated the configuration parser to support loading requests from Python files.
- Added examples and tests for custom request types and plugin functionality.
- Enhanced RequestConfig to accommodate custom request types alongside HTTP and WebSocket.
- Created documentation for the new plugin system and usage examples.
---
 docs/plugin_system.md                         | 238 +++++++++++++++++
 examples/fastapi_app/custom_requests.py       |  92 +++++++
 .../test_config_with_plugins.yaml             |  21 ++
 src/gradual/configs/decorators.py             |  98 +++++++
 src/gradual/configs/parser.py                 |  16 +-
 src/gradual/configs/plugin_loader.py          | 205 ++++++++++++++
 src/gradual/configs/request.py                |  30 ++-
 src/gradual/constants/request_types.py        |   4 +-
 src/gradual/runners/scenario.py               |  57 +++-
 tests/test_custom_request_types.py            | 249 ++++++++++++++++++
 tests/test_plugin_loader.py                   | 163 ++++++++++++
 11 files changed, 1147 insertions(+), 26 deletions(-)
 create mode 100644 docs/plugin_system.md
 create mode 100644 examples/fastapi_app/custom_requests.py
 create mode 100644 examples/fastapi_app/stress_test_configs/test_config_with_plugins.yaml
 create mode 100644 src/gradual/configs/decorators.py
 create mode 100644 src/gradual/configs/plugin_loader.py
 create mode 100644 tests/test_custom_request_types.py
 create mode 100644 tests/test_plugin_loader.py

diff --git a/docs/plugin_system.md b/docs/plugin_system.md
new file mode 100644
index 0000000..10e27eb
--- /dev/null
+++ b/docs/plugin_system.md
@@ -0,0 +1,238 @@
+# Plugin System for Custom Request Functions
+
+The gradual stress testing framework now supports a plugin system that allows you to define custom request functions using Python decorators. This system enables you to create complex, custom request logic while maintaining the framework's configuration-driven approach.
+
+## Overview
+
+The plugin system consists of two main components:
+
+1. **Decorators**: `@request` and `@on_call_completion` decorators for marking functions
+2. **Plugin Loader**: Automatic discovery and loading of decorated functions from Python files
+
+## Decorators
+
+### @request Decorator
+
+The `@request` decorator marks a function as a request function that can be executed during stress testing.
+
+```python
+from gradual.configs.decorators import request
+
+@request(
+    name="my_custom_request",
+    url="http://api.example.com/endpoint",
+    http_method="post",
+    params={"key": "value"},
+    expected_response_time=1.5,
+    auth="bearer",
+    type="http"
+)
+def my_request_function():
+    # Your custom request logic here
+    import requests
+    response = requests.post("http://api.example.com/endpoint", json={"key": "value"})
+    return response.json()
+```
+
+**Parameters:**
+
+- `name` (str, optional): Custom name for the request (defaults to function name)
+- `url` (str): Target URL for the request
+- `params` (dict, optional): Parameters to be sent with the request
+- `http_method` (str): HTTP method to use (default: "get")
+- `expected_response_time` (float): Expected response time in seconds (default: 1.0)
+- `auth` (str, optional): Authentication method to use
+- `type` (str, optional): Type of request (http, websocket, etc.)
+
+### @on_call_completion Decorator
+
+The `@on_call_completion` decorator marks a function as a completion callback that will be executed after a specific request function completes.
+
+```python
+from gradual.configs.decorators import on_call_completion
+
+@on_call_completion(name="my_custom_request")
+def my_completion_callback():
+    # This function will be called after my_custom_request completes
+    print("Request completed!")
+    # You can add logic here like:
+    # - Persisting statistics to database
+    # - Logging results
+    # - Triggering follow-up actions
+```
+
+**Parameters:**
+
+- `name` (str): Name of the request function this callback is associated with
+
+## Using Python Files in Configuration
+
+To use custom request functions in your stress test configuration, you can reference a Python file using the `FROM_REQUEST_YAML_FILE` key:
+
+```yaml
+runs:
+  name: "Custom Requests Test"
+  phases:
+    "phase1":
+      scenarios:
+        "custom_scenario":
+          requests: "FROM_REQUEST_YAML_FILE"
+          request_file: "path/to/your/requests.py"
+          min_concurrency: 1
+          max_concurrency: 10
+          ramp_up_multiply: [1, 2, 3]
+          ramp_up_wait: [1, 1, 1]
+          iterate_through_requests: true
+      run_time: 60
+```
+
+The framework will automatically:
+
+1. Load the Python file
+2. Discover all functions decorated with `@request`
+3. Create `RequestConfig` objects from these functions
+4. Associate completion callbacks with their respective request functions
+
+## Example: Complete Custom Request File
+
+Here's a complete example of a Python file with custom request functions:
+
+```python
+# custom_requests.py
+import time
+import requests
+from gradual.configs.decorators import request, on_call_completion
+
+@request(
+    name="get_user_data",
+    url="http://api.example.com/users",
+    http_method="get",
+    expected_response_time=1.0
+)
+def fetch_user_data():
+    """Fetch user data from API."""
+    response = requests.get("http://api.example.com/users")
+    return response.json()
+
+@on_call_completion(name="get_user_data")
+def user_data_completed():
+    """Callback when user data request completes."""
+    print("User data request completed")
+    # Add your completion logic here
+
+@request(
+    name="create_user",
+    url="http://api.example.com/users",
+    http_method="post",
+    params={"name": "Test User", "email": "test@example.com"},
+    expected_response_time=2.0
+)
+def create_new_user():
+    """Create a new user."""
+    data = {"name": "Test User", "email": "test@example.com"}
+    response = requests.post("http://api.example.com/users", json=data)
+    return response.json()
+
+@on_call_completion(name="create_user")
+def user_created():
+    """Callback when user creation completes."""
+    print("User creation completed")
+
+@request(
+    name="websocket_connection",
+    url="ws://api.example.com/ws",
+    expected_response_time=0.5,
+    type="websocket"
+)
+def connect_websocket():
+    """Establish WebSocket connection."""
+    # WebSocket implementation would go here
+    time.sleep(0.1)  # Simulate connection time
+    return {"status": "connected"}
+
+@on_call_completion(name="websocket_connection")
+def websocket_connected():
+    """Callback when WebSocket connects."""
+    print("WebSocket connection established")
+```
+
+## Configuration File Example
+
+```yaml
+# test_config.yaml
+runs:
+  name: "Custom API Test"
+  wait_between_phases: 5
+  phases:
+    "api_testing":
+      scenarios:
+        "user_operations":
+          requests: "FROM_REQUEST_YAML_FILE"
+          request_file: "custom_requests.py"
+          min_concurrency: 1
+          max_concurrency: 5
+          ramp_up_multiply: [1, 2, 3]
+          ramp_up_wait: [2, 2, 2]
+          iterate_through_requests: true
+      run_time: 30
+```
+
+## Advanced Usage
+
+### Custom Logic Without HTTP Requests
+
+You can create request functions that don't make HTTP requests but perform custom logic:
+
+```python
+@request(
+    name="data_processing",
+    expected_response_time=0.5
+)
+def process_data():
+    """Process data without making HTTP requests."""
+    # Your custom processing logic here
+    result = {"processed": True, "timestamp": time.time()}
+    return result
+
+@on_call_completion(name="data_processing")
+def processing_completed():
+    """Callback when data processing completes."""
+    print("Data processing completed")
+```
+
+### Dynamic Parameters
+
+You can use dynamic parameters in your request functions:
+
+```python
+@request(
+    name="dynamic_request",
+    url="http://api.example.com/dynamic",
+    expected_response_time=1.0
+)
+def dynamic_request():
+    """Request with dynamic parameters."""
+    # Generate dynamic parameters
+    timestamp = int(time.time())
+    params = {"timestamp": timestamp, "random": random.randint(1, 100)}
+
+    response = requests.get("http://api.example.com/dynamic", params=params)
+    return response.json()
+```
+
+## Best Practices
+
+1. **Function Names**: Use descriptive function names that indicate what the request does
+2. **Error Handling**: Include proper error handling in your request functions
+3. **Completion Callbacks**: Use completion callbacks for logging, statistics, or cleanup
+4. **Documentation**: Add docstrings to your request functions for clarity
+5. **Modular Design**: Keep related requests in the same file for better organization
+
+## Integration with Existing Framework
+
+The plugin system integrates seamlessly with the existing gradual framework:
+
+- RequestConfig objects created from decorated functions have the same structure as YAML-defined requests
+- Completion callbacks are stored in the request context and can be executed by the framework
+- All existing framework features (concurrency, ramp-up, timing) work with custom request functions
+- The framework handles execution, timing, and statistics collection automatically
diff --git a/examples/fastapi_app/custom_requests.py b/examples/fastapi_app/custom_requests.py
new file mode 100644
index 0000000..d3e4efa
--- /dev/null
+++ b/examples/fastapi_app/custom_requests.py
@@ -0,0 +1,92 @@
+"""
+Example custom request functions using the @request and @on_call_completion decorators.
+
+This file demonstrates how to create custom request functions that can be loaded
+by the gradual stress testing framework.
+"""
+
+import time
+
+import requests
+
+from gradual.configs.decorators import on_call_completion, request
+
+
+@request(
+    name="custom_get_request",
+    url="http://localhost:8000/api/items",
+    http_method="get",
+    expected_response_time=1.0,
+)
+def get_items():
+    """Custom GET request to fetch items."""
+    response = requests.get("http://localhost:8000/api/items")
+    return response.json()
+
+
+@on_call_completion(name="custom_get_request")
+def get_items_complete():
+    """Completion callback for get_items request."""
+    print("GET items request completed")
+    # Here you could persist stats to database
+    # persist_stats()
+
+
+@request(
+    name="custom_post_request",
+    url="http://localhost:8000/api/items",
+    http_method="post",
+    params={"name": "test_item", "description": "Test item for stress testing"},
+    expected_response_time=2.0,
+)
+def create_item():
+    """Custom POST request to create an item."""
+    data = {"name": "test_item", "description": "Test item for stress testing"}
+    response = requests.post("http://localhost:8000/api/items", json=data)
+    return response.json()
+
+
+@on_call_completion(name="custom_post_request")
+def create_item_complete():
+    """Completion callback for create_item request."""
+    print("POST create item request completed")
+    # Here you could persist stats to database
+    # persist_stats()
+
+
+@request(
+    name="custom_websocket_request",
+    url="ws://localhost:8000/ws",
+    expected_response_time=0.5,
+    type="websocket",
+)
+def websocket_connection():
+    """Custom WebSocket connection request."""
+    # This would be implemented with a WebSocket client
+    # For now, just simulate a connection
+    time.sleep(0.1)
+    return {"status": "connected"}
+
+
+@on_call_completion(name="custom_websocket_request")
+def websocket_complete():
+    """Completion callback for websocket_connection request."""
+    print("WebSocket connection request completed")
+    # Here you could persist stats to database
+    # persist_stats()
+
+
+@request(name="simple_function", expected_response_time=0.1)
+def simple_function():
+    """A simple function without HTTP request."""
+    # This could be any custom logic
+    result = {"message": "Hello from custom function", "timestamp": time.time()}
+    return result
+
+
+@on_call_completion(name="simple_function")
+def simple_function_complete():
+    """Completion callback for simple_function."""
+    print("Simple function completed")
+    # Here you could persist stats to database
+    # persist_stats()
diff --git a/examples/fastapi_app/stress_test_configs/test_config_with_plugins.yaml b/examples/fastapi_app/stress_test_configs/test_config_with_plugins.yaml
new file mode 100644
index 0000000..f551fb4
--- /dev/null
+++ b/examples/fastapi_app/stress_test_configs/test_config_with_plugins.yaml
@@ -0,0 +1,21 @@
+runs:
+  name: "Plugin Test Run"
+  wait_between_phases: 5
+  phases:
+    "phase1":
+      scenarios:
+        "custom_requests":
+          requests: "FROM_REQUEST_YAML_FILE"
+          request_file: "custom_requests.py"
+          min_concurrency: 1
+          max_concurrency: 5
+          ramp_up_multiply:
+            - 1
+            - 2
+            - 3
+          ramp_up_wait:
+            - 1
+            - 1
+            - 1
+          iterate_through_requests: true
+      run_time: 30
diff --git a/src/gradual/configs/decorators.py b/src/gradual/configs/decorators.py
new file mode 100644
index 0000000..585f1d7
--- /dev/null
+++ b/src/gradual/configs/decorators.py
@@ -0,0 +1,98 @@
+"""
+The decorators module provides decorators for marking functions as request functions
+and completion callbacks in the stress testing framework.
+
+This module includes:
+1. @request decorator for marking functions as request functions
+2. @on_call_completion decorator for marking completion callbacks
+"""
+
+from functools import wraps
+from typing import Any, Callable, Dict, Optional, cast
+
+
+def request(
+    name: Optional[str] = None,
+    url: str = "",
+    params: Optional[Dict[str, Any]] = None,
+    http_method: str = "get",
+    expected_response_time: float = 1.0,
+    auth: Optional[str] = None,
+    type: Optional[str] = None,
+):
+    """
+    Decorator to mark a function as a request function.
+
+    This decorator marks a function as a request function that can be executed
+    during stress testing. The function will be discovered by the plugin loader
+    and converted into a RequestConfig object.
+
+    Args:
+        name (Optional[str]): Custom name for the request (defaults to function name)
+        url (str): Target URL for the request
+        params (Optional[Dict[str, Any]]): Parameters to be sent with the request
+        http_method (str): HTTP method to use (for HTTP requests)
+        expected_response_time (float): Expected response time in seconds
+        auth (Optional[str]): Authentication method to use
+        type (Optional[str]): Type of request (http, websocket, etc.)
+
+    Returns:
+        Callable: Decorated function with request metadata
+    """
+
+    def decorator(func: Callable) -> Callable:
+        # Store metadata on the function
+        cast(Any, func)._is_request_function = True
+        cast(Any, func)._request_metadata = {
+            "name": name or func.__name__,
+            "url": url,
+            "params": params or {},
+            "http_method": http_method,
+            "expected_response_time": expected_response_time,
+            "auth": auth,
+            "type": type,
+        }
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        # Copy metadata to wrapper
+        cast(Any, wrapper)._is_request_function = True
+        cast(Any, wrapper)._request_metadata = cast(Any, func)._request_metadata
+
+        return wrapper
+
+    return decorator
+
+
+def on_call_completion(name: str):
+    """
+    Decorator to mark a function as a completion callback for a request function.
+
+    This decorator marks a function as a completion callback that will be called
+    after a specific request function completes execution.
+
+    Args:
+        name (str): Name of the request function this callback is associated with
+
+    Returns:
+        Callable: Decorated function with completion callback metadata
+    """
+
+    def decorator(func: Callable) -> Callable:
+        # Store metadata on the function
+        cast(Any, func)._is_completion_callback = True
+        cast(Any, func)._target_request_function = name
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+        # Copy metadata to wrapper
+        cast(Any, wrapper)._is_completion_callback = True
+        cast(Any, wrapper)._target_request_function = name
+
+        return wrapper
+
+    return decorator
diff --git a/src/gradual/configs/parser.py b/src/gradual/configs/parser.py
index 0a3242e..105e943 100644
--- a/src/gradual/configs/parser.py
+++ b/src/gradual/configs/parser.py
@@ -10,9 +10,10 @@
 
 import yaml
 
+from gradual.configs.phase import PhaseConfig
+from gradual.configs.plugin_loader import load_request_configs_from_file
 from gradual.configs.request import RequestConfig
 from gradual.configs.scenario import ScenarioConfig
-from gradual.configs.phase import PhaseConfig
 from gradual.configs.validate import assert_not_empty, validate_min_concurrency
 
 
@@ -160,9 +161,16 @@ def read_configs(self):
                 request_configs = []
 
                 if scenario_data["requests"] == "FROM_REQUEST_YAML_FILE":
-                    request_configs = self.read_request_file(
-                        scenario_data["request_file"]
-                    )
+                    request_file_path = scenario_data["request_file"]
+                    # Check if it's a Python file
+                    if request_file_path.endswith(".py"):
+                        request_configs = load_request_configs_from_file(
+                            request_file_path
+                        )
+                    else:
+                        request_configs = self.read_request_file(
+                            Path(request_file_path)
+                        )
                 else:
                     for scenario_request_name in scenario_data["requests"]:
                         request = params_config["requests"][scenario_request_name]
diff --git a/src/gradual/configs/plugin_loader.py b/src/gradual/configs/plugin_loader.py
new file mode 100644
index 0000000..2b6b0f8
--- /dev/null
+++ b/src/gradual/configs/plugin_loader.py
@@ -0,0 +1,205 @@
+"""
+The plugin_loader module provides functionality to load custom request functions
+from Python files and convert them into RequestConfig objects for stress testing.
+
+This module handles:
+1. Importing Python files containing decorated request functions
+2. Discovering functions decorated with @request
+3. Creating RequestConfig objects from these functions
+4. Managing completion callbacks for request functions
+"""
+
+import importlib.util
+import inspect
+import sys
+from pathlib import Path
+from typing import Any, Callable, Dict, List, Optional
+
+from gradual.configs.request import RequestConfig
+from gradual.constants.request_types import RequestType
+
+
+class PluginLoader:
+    """
+    Loader for custom request functions from Python files.
+
+    This class handles the discovery and loading of request functions that are
+    decorated with @request and their associated completion callbacks.
+    """
+
+    def __init__(self):
+        self._request_functions: Dict[str, Callable] = {}
+        self._completion_callbacks: Dict[str, Callable] = {}
+        self._loaded_modules: Dict[str, Any] = {}
+
+    def load_plugin_file(self, filepath: str) -> List[RequestConfig]:
+        """
+        Load a Python file and extract RequestConfig objects from decorated functions.
+
+        Args:
+            file_path (str): Path to the Python file containing request functions
+
+        Returns:
+            List[RequestConfig]: List of RequestConfig objects created from decorated functions
+
+        Raises:
+            FileNotFoundError: If the specified file doesn't exist
+            ImportError: If there's an error importing the module
+        """
+        file_path = Path(filepath)
+        if not file_path.exists():
+            raise FileNotFoundError(f"Plugin file not found: {file_path}")
+
+        # Load the module
+        module_name = file_path.stem
+        spec = importlib.util.spec_from_file_location(module_name, file_path)
+        if spec is None or spec.loader is None:
+            raise ImportError(f"Could not load module from {file_path}")
+
+        module = importlib.util.module_from_spec(spec)
+        sys.modules[module_name] = module
+        spec.loader.exec_module(module)
+
+        # Store the loaded module
+        self._loaded_modules[module_name] = module
+
+        # Discover request functions and completion callbacks
+        self._discover_functions(module)
+
+        # Create RequestConfig objects
+        return self._create_request_configs()
+
+    def _discover_functions(self, module: Any) -> None:
+        """
+        Discover request functions and completion callbacks in a module.
+
+        Args:
+            module: The loaded module to inspect
+        """
+        for name, obj in inspect.getmembers(module):
+            if inspect.isfunction(obj):
+                # Check if function has request decorator metadata
+                if hasattr(obj, "_is_request_function"):
+                    self._request_functions[name] = obj
+
+                # Check if function has completion callback metadata
+                if hasattr(obj, "_is_completion_callback"):
+                    target_function = getattr(obj, "_target_request_function", None)
+                    if target_function:
+                        self._completion_callbacks[target_function] = obj
+
+    def _create_request_configs(self) -> List[RequestConfig]:
+        """
+        Create RequestConfig objects from discovered request functions.
+
+        Returns:
+            List[RequestConfig]: List of RequestConfig objects
+        """
+        configs = []
+
+        for func_name, func in self._request_functions.items():
+            # Get function metadata
+            func_metadata = getattr(func, "_request_metadata", {})
+
+            # Use the name from metadata if provided, otherwise use function name
+            request_name = func_metadata.get("name", func_name)
+
+            # Map string type to RequestType enum
+            type_str = func_metadata.get("type", None)
+            # Will trigger auto-detection in RequestConfig
+            request_type = None
+            if isinstance(type_str, str):
+                if type_str.lower() == "http":
+                    request_type = RequestType.http
+                elif type_str.lower() == "websocket":
+                    request_type = RequestType.websocket
+                else:
+                    # Default for unknown or "custom"
+                    request_type = RequestType.custom
+
+            # Create RequestConfig
+            config = RequestConfig(
+                name=request_name,
+                params=func_metadata.get("params", {}),
+                http_method=func_metadata.get("http_method", "get"),
+                expected_response_time=func_metadata.get("expected_response_time", 1.0),
+                context={
+                    "function": func,
+                    "completion_callback": self._completion_callbacks.get(request_name),
+                    "metadata": func_metadata,
+                },
+                url=func_metadata.get("url", ""),
+                auth=func_metadata.get("auth", None),
+                type=request_type,
+            )
+
+            configs.append(config)
+
+        return configs
+
+    def get_request_function(self, name: str) -> Optional[Callable]:
+        """
+        Get a request function by name.
+
+        Args:
+            name (str): Name of the request function
+
+        Returns:
+            Optional[Callable]: The request function if found, None otherwise
+        """
+        return self._request_functions.get(name)
+
+    def get_completion_callback(self, request_name: str) -> Optional[Callable]:
+        """
+        Get a completion callback for a request function.
+
+        Args:
+            request_name (str): Name of the request function
+
+        Returns:
+            Optional[Callable]: The completion callback if found, None otherwise
+        """
+        return self._completion_callbacks.get(request_name)
+
+    def clear_state(self):
+        """
+        Clear the internal state of the plugin loader.
+
+        This method clears all discovered functions and callbacks, useful for
+        testing or when loading multiple files that should be independent.
+        """
+        self._request_functions.clear()
+        self._completion_callbacks.clear()
+        self._loaded_modules.clear()
+
+
+# Global plugin loader instance
+_plugin_loader = PluginLoader()
+
+
+def get_plugin_loader() -> PluginLoader:
+    """
+    Get the global plugin loader instance.
+
+    Returns:
+        PluginLoader: The global plugin loader instance
+    """
+    return _plugin_loader
+
+
+def load_request_configs_from_file(file_path: str) -> List[RequestConfig]:
+    """
+    Load RequestConfig objects from a Python file.
+
+    This is a convenience function that uses the global plugin loader
+    to load request configurations from a Python file.
+
+    Args:
+        file_path (str): Path to the Python file containing request functions
+
+    Returns:
+        List[RequestConfig]: List of RequestConfig objects
+    """
+    # Clear state before loading to ensure clean slate
+    _plugin_loader.clear_state()
+    return _plugin_loader.load_plugin_file(file_path)
diff --git a/src/gradual/configs/request.py b/src/gradual/configs/request.py
index 0cc8d7e..f35e3dc 100644
--- a/src/gradual/configs/request.py
+++ b/src/gradual/configs/request.py
@@ -1,7 +1,7 @@
 """
 The request module provides configuration classes and utilities for managing API request
 configurations in the stress testing framework. It includes support for different
-request types (HTTP, WebSocket) and their specific parameters.
+request types (HTTP, WebSocket, Custom) and their specific parameters.
 """
 
 from dataclasses import dataclass
@@ -34,6 +34,7 @@ def check_websocket_or_http(url):
         return RequestType.http
     if url_type in RequestType.websocket.value:
         return RequestType.websocket
+    return None
 
 
 @dataclass
@@ -42,8 +43,8 @@ class RequestConfig:
     Configuration class for API requests in stress testing.
 
     This class defines the structure and parameters for individual API requests,
-    supporting both HTTP and WebSocket protocols. It includes validation and
-    automatic type detection based on the URL.
+    supporting HTTP, WebSocket, and custom protocols. It includes validation and
+    automatic type detection based on the URL, while respecting explicitly set types.
 
     Attributes:
         name (str): Unique identifier for this request configuration
@@ -53,7 +54,7 @@ class RequestConfig:
         context (Optional[dict[str, Any]]): Additional context for the request
         url (str): Target URL for the request
         auth (Optional[str]): Authentication method to use
-        type (Optional[RequestType]): Type of request (HTTP or WebSocket)
+        type (Optional[RequestType]): Type of request (HTTP, WebSocket, or Custom)
     """
 
     name: str
@@ -63,13 +64,24 @@ class RequestConfig:
     context: Optional[dict[str, Any]] = None
     url: str = ""
     auth: Optional[str] = None
-    type: Optional[RequestType] = RequestType.http
+    type: Optional[RequestType] = None
 
     def __post_init__(self):
         """
         Post-initialization hook to set the request type based on the URL.
-
-        This method is automatically called after initialization to determine
-        the appropriate request type based on the URL protocol.
+        Only auto-detect if no explicit type was provided or if it's the default HTTP type.
         """
-        self.type = check_websocket_or_http(self.url)
+        if self.context is None:
+            self.context = {}
+
+        if self.type is None:
+            # Auto-detect type based on URL
+            detected_type = check_websocket_or_http(self.url)
+            if detected_type:
+                self.type = detected_type
+            else:
+                # If no URL or unrecognized protocol, default to custom
+                self.type = RequestType.custom
+        elif self.type == RequestType.http and not self.url:
+            # If explicitly set to HTTP but no URL, default to custom
+            self.type = RequestType.custom
diff --git a/src/gradual/constants/request_types.py b/src/gradual/constants/request_types.py
index fcc2cbd..049dc0c 100644
--- a/src/gradual/constants/request_types.py
+++ b/src/gradual/constants/request_types.py
@@ -1,7 +1,7 @@
 """
 The request_types module provides the RequestType enum which defines the supported
 types of API requests in the stress testing framework. It includes HTTP and WebSocket
-protocols with their respective URL schemes.
+protocols with their respective URL schemes, plus support for custom request types.
 """
 
 from enum import Enum
@@ -18,7 +18,9 @@ class RequestType(Enum):
     Attributes:
         websocket (list[str]): List of WebSocket URL schemes (wss, ws)
         http (list[str]): List of HTTP URL schemes (http, https)
+        custom (list[str]): Custom request types (empty list for any custom type)
     """
 
     websocket = ["wss", "ws"]
     http = ["http", "https"]
+    custom: list[str] = []  # For custom request types that don't follow URL patterns
diff --git a/src/gradual/runners/scenario.py b/src/gradual/runners/scenario.py
index a1fb0b0..7697228 100644
--- a/src/gradual/runners/scenario.py
+++ b/src/gradual/runners/scenario.py
@@ -4,15 +4,17 @@
 """
 
 from logging import info
+
+import gevent
+from tabulate import tabulate
+
 from gradual.configs.scenario import ScenarioConfig
 from gradual.constants.request_types import RequestType
+from gradual.runners.iterators import RequestIterator
 from gradual.runners.request.base import _Request
 from gradual.runners.request.Http import HttpRequest
 from gradual.runners.request.SocketIO import SocketRequest
-from gradual.runners.iterators import RequestIterator
 from gradual.runners.session import HTTPSession
-import gevent
-from tabulate import tabulate
 
 
 class Scenario:
@@ -103,11 +105,35 @@ def do_ramp_up(self, ramp_up_value):
                     iterator=iterator,
                 )
             else:
-                request = _Request(
-                    scenario_name=self.scenario_config.name,
-                    run_once=self.scenario_config.run_once,
-                    iterator=iterator,
+                # Create a dynamic _Request subclass with the user's run function
+                custom_function = current_request_type.context.get("function")
+                completion_callback = current_request_type.context.get(
+                    "completion_callback"
                 )
+
+                if custom_function:
+                    # Create a dynamic class that inherits from _Request
+                    class CustomRequestClass(_Request):
+                        def run(self):
+                            return custom_function()  # noqa: B023
+
+                        def on_request_completion(self, *args, **kwargs):
+                            if completion_callback:  # noqa: B023
+                                completion_callback()  # noqa: B023
+
+                    request = CustomRequestClass(
+                        scenario_name=self.scenario_config.name,
+                        run_once=self.scenario_config.run_once,
+                        iterator=iterator,
+                    )
+
+                else:
+                    request = _Request(
+                        scenario_name=self.scenario_config.name,
+                        run_once=self.scenario_config.run_once,
+                        iterator=iterator,
+                    )
+
             self.running_request_tasks.append(gevent.spawn(request.run))
             current_concurrency += 1
             self.requests.append(request)
@@ -130,7 +156,8 @@ def execute(self):
         5. Provides detailed logging of execution progress
         """
         info(
-            f"Starting the testiung with minimum concurrency i.e., {self.scenario_config.min_concurrency}, scenario: {self.scenario_config.name}"
+            f"Starting the testiung with minimum concurrency i.e., "
+            f"{self.scenario_config.min_concurrency}, scenario: {self.scenario_config.name}"
         )
 
         # Current index of ramp up and ramp up wait array.
@@ -153,16 +180,22 @@ def execute(self):
 
             # Calculating by how much we have to ramp up in this iteration.
             if self.scenario_config.multiply:
-                # Suppose we want to ramp up the total requests by 2x and there are already x requests running in an infinite loop.
-                # Then, total requests need to be added is 2x = already_running_request(x) * (multiplication_facotr(2) -1 ) to make the concurrency 2x.
+                # Suppose we want to ramp up the total requests by 2x and
+                # there are already x requests running in an infinite loop.
+                # Then, total requests need to be added is
+                # 2x = already_running_request(x) * (multiplication_facotr(2) -1 )
+                # to make the concurrency 2x.
                 if not self.scenario_config.run_once:
                     ramp_up_val = len(self.running_request_tasks) * (
                         self.scenario_config.ramp_up[ramp_up_idx] - 1
                     )
 
-                # Suppose we want to ramp up the total requests by 2x and there are already x requests with run_once True.
+                # Suppose we want to ramp up the total requests by 2x and
+                # there are already x requests with run_once True.
                 # That means we are ramping up after the requests are completed.
-                # Then, total requests needs to be added is 2x = already_running_request(x) * (multiplication_facotr(2)) to make the concurrency 2x.
+                # Then, total requests needs to be added is
+                # 2x = already_running_request(x) * (multiplication_facotr(2))
+                # to make the concurrency 2x.
                 else:
                     ramp_up_val = len(self.running_request_tasks) * (
                         self.scenario_config.ramp_up[ramp_up_idx]
diff --git a/tests/test_custom_request_types.py b/tests/test_custom_request_types.py
new file mode 100644
index 0000000..0a49a4d
--- /dev/null
+++ b/tests/test_custom_request_types.py
@@ -0,0 +1,249 @@
+"""
+Tests for custom request types functionality.
+"""
+
+import os
+import tempfile
+
+from gradual.configs.decorators import request
+from gradual.configs.plugin_loader import load_request_configs_from_file
+from gradual.configs.request import RequestConfig
+from gradual.constants.request_types import RequestType
+
+
+def test_custom_request_type_without_url():
+    """Test that custom request types work without URLs."""
+
+    # Create a temporary Python file with custom request functions
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            '''
+from gradual.configs.decorators import request, on_call_completion
+
+@request(
+    name="database_query",
+    type="custom",
+    expected_response_time=0.5
+)
+def query_database():
+    """Custom database query operation."""
+    return {"result": "data", "rows": 100}
+
+@on_call_completion(name="database_query")
+def db_completion():
+    return "Database query completed"
+
+@request(
+    name="file_operation",
+    type="custom",
+    expected_response_time=1.0
+)
+def process_file():
+    """Custom file processing operation."""
+    return {"status": "processed", "files": 5}
+
+@on_call_completion(name="file_operation")
+def file_completion():
+    return "File processing completed"
+'''
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that RequestConfig objects were created
+        assert len(configs) == 2
+
+        # Check database query config
+        db_config = next(c for c in configs if c.name == "database_query")
+        assert db_config.type == RequestType.custom
+        assert db_config.url == ""
+        assert db_config.expected_response_time == 0.5
+        assert db_config.context["function"].__name__ == "query_database"
+        assert db_config.context["completion_callback"].__name__ == "db_completion"
+
+        # Check file operation config
+        file_config = next(c for c in configs if c.name == "file_operation")
+        assert file_config.type == RequestType.custom
+        assert file_config.url == ""
+        assert file_config.expected_response_time == 1.0
+        assert file_config.context["function"].__name__ == "process_file"
+        assert file_config.context["completion_callback"].__name__ == "file_completion"
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_mixed_request_types():
+    """Test mixing HTTP, WebSocket, and custom request types."""
+
+    # Create a temporary Python file with mixed request types
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+from gradual.configs.decorators import request, on_call_completion
+
+@request(
+    name="http_request",
+    url="http://api.example.com/users",
+    type="http",
+    http_method="get",
+    expected_response_time=1.0
+)
+def http_request():
+    return {"users": []}
+
+@request(
+    name="websocket_request",
+    url="ws://api.example.com/ws",
+    type="websocket",
+    expected_response_time=0.5
+)
+def websocket_request():
+    return {"connected": True}
+
+@request(
+    name="custom_request",
+    type="custom",
+    expected_response_time=0.3
+)
+def custom_request():
+    return {"custom": "result"}
+
+@on_call_completion(name="http_request")
+def http_completion():
+    return "HTTP request completed"
+
+@on_call_completion(name="websocket_request")
+def ws_completion():
+    return "WebSocket request completed"
+
+@on_call_completion(name="custom_request")
+def custom_completion():
+    return "Custom request completed"
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that all RequestConfig objects were created
+        assert len(configs) == 3
+
+        # Check HTTP request
+        http_config = next(c for c in configs if c.name == "http_request")
+        assert http_config.type == RequestType.http
+        assert http_config.url == "http://api.example.com/users"
+        assert http_config.http_method == "get"
+
+        # Check WebSocket request
+        ws_config = next(c for c in configs if c.name == "websocket_request")
+        assert ws_config.type == RequestType.websocket
+        assert ws_config.url == "ws://api.example.com/ws"
+
+        # Check custom request
+        custom_config = next(c for c in configs if c.name == "custom_request")
+        assert custom_config.type == RequestType.custom
+        assert custom_config.url == ""
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_auto_detection_with_custom_override():
+    """Test that explicit type overrides auto-detection."""
+
+    # Create a temporary Python file with custom type override
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+from gradual.configs.decorators import request
+
+@request(
+    name="custom_http",
+    url="http://api.example.com/data",
+    type="custom",  # Explicitly set to custom despite HTTP URL
+    expected_response_time=1.0
+)
+def custom_http_request():
+    return {"custom": "http-like"}
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that the type is custom despite HTTP URL
+        assert len(configs) == 1
+        config = configs[0]
+        assert config.type == RequestType.custom
+        assert config.url == "http://api.example.com/data"
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_unknown_type_defaults_to_custom():
+    """Test that unknown types default to custom."""
+
+    # Create a temporary Python file with unknown type
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+from gradual.configs.decorators import request
+
+@request(
+    name="unknown_type",
+    type="unknown_type",  # Unknown type
+    expected_response_time=1.0
+)
+def unknown_request():
+    return {"unknown": "type"}
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that unknown type defaults to custom
+        assert len(configs) == 1
+        config = configs[0]
+        assert config.type == RequestType.custom
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_custom_request_execution():
+    """Test that custom request functions can be executed."""
+
+    @request(name="test_custom", type="custom", expected_response_time=0.1)
+    def test_function():
+        return {"message": "Hello from custom function"}
+
+    # Create RequestConfig manually to test execution
+    config = RequestConfig(
+        name="test_custom",
+        params={},
+        http_method="get",
+        expected_response_time=0.1,
+        context={"function": test_function},
+        type=RequestType.custom,
+    )
+
+    # Execute the function
+    result = config.context["function"]()
+    assert result == {"message": "Hello from custom function"}
+    assert config.type == RequestType.custom
diff --git a/tests/test_plugin_loader.py b/tests/test_plugin_loader.py
new file mode 100644
index 0000000..b24e22c
--- /dev/null
+++ b/tests/test_plugin_loader.py
@@ -0,0 +1,163 @@
+"""
+Tests for the plugin loader and decorator functionality.
+"""
+
+import os
+import tempfile
+
+import pytest
+
+from gradual.configs.decorators import on_call_completion, request
+from gradual.configs.plugin_loader import load_request_configs_from_file
+from gradual.configs.request import RequestConfig
+
+
+def test_request_decorator():
+    """Test that the @request decorator properly marks functions."""
+
+    @request(name="test_request", expected_response_time=1.0)
+    def test_func():
+        return "test"
+
+    assert hasattr(test_func, "_is_request_function")
+    assert test_func._is_request_function is True
+    assert hasattr(test_func, "_request_metadata")
+    assert test_func._request_metadata["name"] == "test_request"
+    assert test_func._request_metadata["expected_response_time"] == 1.0
+
+
+def test_on_call_completion_decorator():
+    """Test that the @on_call_completion decorator properly marks functions."""
+
+    @on_call_completion(name="test_request")
+    def completion_func():
+        return "completed"
+
+    assert hasattr(completion_func, "_is_completion_callback")
+    assert completion_func._is_completion_callback is True
+    assert hasattr(completion_func, "_target_request_function")
+    assert completion_func._target_request_function == "test_request"
+
+
+def test_plugin_loader_discovery():
+    """Test that the plugin loader can discover decorated functions."""
+
+    # Create a temporary Python file with decorated functions
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+from gradual.configs.decorators import request, on_call_completion
+
+@request(name="test_request", expected_response_time=1.0)
+def test_func():
+    return "test"
+
+@on_call_completion(name="test_request")
+def completion_func():
+    return "completed"
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that RequestConfig objects were created
+        assert len(configs) == 1
+        config = configs[0]
+        assert isinstance(config, RequestConfig)
+        assert config.name == "test_request"
+        assert config.expected_response_time == 1.0
+
+        # Verify that the function and callback are stored in context
+        assert "function" in config.context
+        assert "completion_callback" in config.context
+        assert config.context["function"].__name__ == "test_func"
+        assert config.context["completion_callback"].__name__ == "completion_func"
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_plugin_loader_multiple_functions():
+    """Test that the plugin loader can handle multiple decorated functions."""
+
+    # Create a temporary Python file with multiple decorated functions
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+from gradual.configs.decorators import request, on_call_completion
+
+@request(name="request1", expected_response_time=1.0)
+def func1():
+    return "test1"
+
+@request(name="request2", expected_response_time=2.0)
+def func2():
+    return "test2"
+
+@on_call_completion(name="request1")
+def completion1():
+    return "completed1"
+
+@on_call_completion(name="request2")
+def completion2():
+    return "completed2"
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # Load the plugin file
+        configs = load_request_configs_from_file(temp_file)
+
+        # Verify that both RequestConfig objects were created
+        assert len(configs) == 2
+
+        # Find the configs by name
+        config1 = next(c for c in configs if c.name == "request1")
+        config2 = next(c for c in configs if c.name == "request2")
+
+        assert config1.expected_response_time == 1.0
+        assert config2.expected_response_time == 2.0
+
+        # Verify callbacks are correctly associated
+        assert config1.context["completion_callback"].__name__ == "completion1"
+        assert config2.context["completion_callback"].__name__ == "completion2"
+
+    finally:
+        # Clean up
+        os.unlink(temp_file)
+
+
+def test_plugin_loader_file_not_found():
+    """Test that the plugin loader raises appropriate error for missing files."""
+
+    with pytest.raises(FileNotFoundError):
+        load_request_configs_from_file("nonexistent_file.py")
+
+
+def test_plugin_loader_invalid_python():
+    """Test that the plugin loader handles invalid Python files gracefully."""
+
+    # Create a temporary file with invalid Python syntax
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(
+            """
+invalid python syntax here
+@request(name="test")
+def test_func():
+    return "test"
+"""
+        )
+        temp_file = f.name
+
+    try:
+        # This should raise a syntax error
+        with pytest.raises(SyntaxError):
+            load_request_configs_from_file(temp_file)
+    finally:
+        # Clean up
+        os.unlink(temp_file)

From 750f2eaf460aaff3c52e8df9858e5f22fc2cdd96 Mon Sep 17 00:00:00 2001
From: Subham Agrawal <subhamagrawal@Subhams-MacBook-Pro.local>
Date: Fri, 8 Aug 2025 23:53:42 +0530
Subject: [PATCH 2/3] fix: lint errors

---
 src/gradual/configs/plugin_loader.py   | 3 ++-
 src/gradual/constants/request_types.py | 2 +-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/src/gradual/configs/plugin_loader.py b/src/gradual/configs/plugin_loader.py
index 2b6b0f8..649f9c0 100644
--- a/src/gradual/configs/plugin_loader.py
+++ b/src/gradual/configs/plugin_loader.py
@@ -13,6 +13,7 @@
 import inspect
 import sys
 from pathlib import Path
+from types import ModuleType
 from typing import Any, Callable, Dict, List, Optional
 
 from gradual.configs.request import RequestConfig
@@ -69,7 +70,7 @@ def load_plugin_file(self, filepath: str) -> List[RequestConfig]:
         # Create RequestConfig objects
         return self._create_request_configs()
 
-    def _discover_functions(self, module: Any) -> None:
+    def _discover_functions(self, module: ModuleType) -> None:
         """
         Discover request functions and completion callbacks in a module.
 
diff --git a/src/gradual/constants/request_types.py b/src/gradual/constants/request_types.py
index 049dc0c..a4c3815 100644
--- a/src/gradual/constants/request_types.py
+++ b/src/gradual/constants/request_types.py
@@ -23,4 +23,4 @@ class RequestType(Enum):
 
     websocket = ["wss", "ws"]
     http = ["http", "https"]
-    custom: list[str] = []  # For custom request types that don't follow URL patterns
+    custom: list = []  # For custom request types that don't follow URL patterns

From 048019beab95fcaefe0f9d6315e9bc00c33b9ebd Mon Sep 17 00:00:00 2001
From: Subham Agrawal <subhamagrawal@Subhams-MacBook-Pro.local>
Date: Fri, 8 Aug 2025 23:58:52 +0530
Subject: [PATCH 3/3] fix: lint related errors

---
 pyproject.toml                         | 2 +-
 src/gradual/constants/request_types.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index d17e825..2a34038 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -125,7 +125,7 @@ strict_optional = true
 warn_redundant_casts = true
 warn_return_any = true
 warn_unused_ignores = true
-disable_error_code = ["no-untyped-def", "import-untyped", "call-arg"]
+disable_error_code = ["no-untyped-def", "import-untyped", "call-arg", "annotation-unchecked"]
 [[tool.mypy.overrides]]
 module = ["requests_kerberos"]
 ignore_missing_imports = true
diff --git a/src/gradual/constants/request_types.py b/src/gradual/constants/request_types.py
index a4c3815..f8581af 100644
--- a/src/gradual/constants/request_types.py
+++ b/src/gradual/constants/request_types.py
@@ -23,4 +23,4 @@ class RequestType(Enum):
 
     websocket = ["wss", "ws"]
     http = ["http", "https"]
-    custom: list = []  # For custom request types that don't follow URL patterns
+    custom = ["custom"]  # For custom request types that don't follow URL patterns