devdox-ai-locust 0.1.3.post1__tar.gz → 0.1.4__tar.gz

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: devdox_ai_locust
-Version: 0.1.3.post1
+Version: 0.1.4
 Summary: AI-powered Locust load test generator from API documentation
 Author-email: Hayat Bourji <hayat.bourgi@montyholding.com>
 Maintainer-email: Hayat Bourji <hayat.bourgi@montyholding.com>
@@ -73,19 +73,27 @@ Dynamic: license-file
 DevDox AI Locust automatically generates comprehensive Locust load testing scripts from your API documentation (OpenAPI/Swagger specs). Using advanced AI capabilities, it creates realistic test scenarios, handles complex authentication flows, and generates production-ready performance tests.
 
 
-## 🆕 What's New in 0.1.3.post1
+## 🆕 What's New in 0.1.4
 
-### Performance & Reliability Improvements
+### MongoDB Integration
 
-- **🚀 Asynchronous API Calls**: Migrated from `Together` to `AsyncTogether` for non-blocking API interactions, significantly improving generation speed and responsiveness
-- **⚡ Enhanced Timeout Handling**: Implemented robust timeout logic with configurable retry mechanisms for better error resilience
-- **🔧 Improved Code Extraction**: Enhanced `<code>` block parsing with detailed validation, multiple fallback scenarios, and better error messages
-- **🛡️ Better Error Management**: Comprehensive error handling throughout the AI generation pipeline with graceful degradation
+- Added a new data provider class: MongoDataProvider
+- Connects Locust test data generation directly to MongoDB
+- Enables realistic test data retrieval for entities like users, products, orders, affiliates, etc.
+- Supports **real data** from the database and **synthetic fallback generation** when MongoDB is disabled or unavailable
 
-### Bug Fixes & Stability
+#### 2. **New MongoDataProvider Methods**
+| Method                                                   | Description |
+|----------------------------------------------------------|--------------|
+| `get_document(collection_name)`                      | Retrieves a single realistic document from MongoDB or fallback generator |
+| `get_multiple_documents(collection_name, count=10, query=None)` | Retrieves multiple documents or generates them in batches |
+| `clear_cache()`                                          | Clears in-memory cached data for all collections |
+| `get_stats()`                                            | Returns usage and cache statistics for debugging and optimization |
+
+#### 3. **Smart Fallbacks**
+If MongoDB is disabled (`enable_mongodb = false` in `db_config.py`),  
+the system automatically switches to **synthetic data generation** using the LLM-based `TestDataGenerator`.
 
-- Fixed edge cases in code block extraction where malformed responses could cause generation failures
-- Improved retry logic to handle transient API errors without interrupting the generation process
 
 ## ✨ Features
 
@@ -139,14 +147,21 @@ devdox_ai_locust generate --openapi-url https://api.example.com/openapi.json --o
 
 # Generate with custom configuration
 devdox_ai_locust generate \
-   https://petstore.swagger.io/v3/swagger.json \
+   https://petstore3.swagger.io/api/v3/openapi.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   
+ # Generate with db integration 
+devdox_ai_locust generate \
+  https://petstore3.swagger.io/api/v3/openapi.json \
+  --output ./petstore-tests \
+  --db-type mongo \
 ```
 
 
 
+
+
 ## 📖 Documentation
 
 ### Command Line Interface

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/README.md

@@ -9,19 +9,27 @@
 DevDox AI Locust automatically generates comprehensive Locust load testing scripts from your API documentation (OpenAPI/Swagger specs). Using advanced AI capabilities, it creates realistic test scenarios, handles complex authentication flows, and generates production-ready performance tests.
 
 
-## 🆕 What's New in 0.1.3.post1
+## 🆕 What's New in 0.1.4
 
-### Performance & Reliability Improvements
+### MongoDB Integration
 
-- **🚀 Asynchronous API Calls**: Migrated from `Together` to `AsyncTogether` for non-blocking API interactions, significantly improving generation speed and responsiveness
-- **⚡ Enhanced Timeout Handling**: Implemented robust timeout logic with configurable retry mechanisms for better error resilience
-- **🔧 Improved Code Extraction**: Enhanced `<code>` block parsing with detailed validation, multiple fallback scenarios, and better error messages
-- **🛡️ Better Error Management**: Comprehensive error handling throughout the AI generation pipeline with graceful degradation
+- Added a new data provider class: MongoDataProvider
+- Connects Locust test data generation directly to MongoDB
+- Enables realistic test data retrieval for entities like users, products, orders, affiliates, etc.
+- Supports **real data** from the database and **synthetic fallback generation** when MongoDB is disabled or unavailable
 
-### Bug Fixes & Stability
+#### 2. **New MongoDataProvider Methods**
+| Method                                                   | Description |
+|----------------------------------------------------------|--------------|
+| `get_document(collection_name)`                      | Retrieves a single realistic document from MongoDB or fallback generator |
+| `get_multiple_documents(collection_name, count=10, query=None)` | Retrieves multiple documents or generates them in batches |
+| `clear_cache()`                                          | Clears in-memory cached data for all collections |
+| `get_stats()`                                            | Returns usage and cache statistics for debugging and optimization |
+
+#### 3. **Smart Fallbacks**
+If MongoDB is disabled (`enable_mongodb = false` in `db_config.py`),  
+the system automatically switches to **synthetic data generation** using the LLM-based `TestDataGenerator`.
 
-- Fixed edge cases in code block extraction where malformed responses could cause generation failures
-- Improved retry logic to handle transient API errors without interrupting the generation process
 
 ## ✨ Features
 
@@ -75,14 +83,21 @@ devdox_ai_locust generate --openapi-url https://api.example.com/openapi.json --o
 
 # Generate with custom configuration
 devdox_ai_locust generate \
-   https://petstore.swagger.io/v3/swagger.json \
+   https://petstore3.swagger.io/api/v3/openapi.json \
   --output ./petstore-tests \
   --together-api-key your_api_key \
   
+ # Generate with db integration 
+devdox_ai_locust generate \
+  https://petstore3.swagger.io/api/v3/openapi.json \
+  --output ./petstore-tests \
+  --db-type mongo \
 ```
 
 
 
+
+
 ## 📖 Documentation
 
 ### Command Line Interface

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "devdox_ai_locust"
-version = "0.1.3.post1"
+version = "0.1.4"
 description = "AI-powered Locust load test generator from API documentation"
 readme = "README.md"
 license = {text = "Apache-2.0" }
@@ -101,7 +101,7 @@ exclude = ["tests*"]
 
 
 [tool.setuptools.package-data]
-devdox_ai_locust = ["schemas/*.json", "templates/*.j2","prompt/*.j2", "py.typed","*.j2"]
+devdox_ai_locust = ["schemas/*.json", "templates/*.j2","templates/mongo/*","prompt/*.j2", "py.typed","*.j2"]
 
 # flake8 configuration (legacy - remove if using ruff)
 [tool.flake8]

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/src/devdox_ai_locust/cli.py

@@ -197,19 +197,20 @@ async def _generate_and_create_tests(
     custom_requirement: Optional[str] = "",
     host: Optional[str] = "0.0.0.0",
     auth: bool = False,
+    db_type: str = "",
 ) -> List[Dict[Any, Any]]:
     """Generate tests using AI and create test files"""
     together_client = AsyncTogether(api_key=api_key)
 
     with console.status("[bold green]Generating Locust tests with AI..."):
         generator = HybridLocustGenerator(ai_client=together_client)
-
         test_files, test_directories = await generator.generate_from_endpoints(
             endpoints=endpoints,
             api_info=api_info,
             custom_requirement=custom_requirement,
             target_host=host,
             include_auth=auth,
+            db_type=db_type,
         )
 
     # Create test files
@@ -271,6 +272,12 @@ def cli(ctx: click.Context, verbose: bool) -> None:
 )
 @click.option("--host", "-H", type=str, help="Target host URL")
 @click.option("--auth/--no-auth", default=True, help="Include authentication in tests")
+@click.option(
+    "--db-type",
+    type=click.Choice(["", "mongo", "postgresql"], case_sensitive=False),
+    default="",
+    help="Database type for testing (empty for no database, mongo, or postgresql)",
+)
 @click.option("--dry-run", is_flag=True, help="Generate tests without running them")
 @click.option(
     "--custom-requirement", type=str, help="Custom requirements for test generation"
@@ -291,6 +298,7 @@ def generate(
     run_time: str,
     host: Optional[str],
     auth: bool,
+    db_type: str,
     dry_run: bool,
     custom_requirement: Optional[str],
     together_api_key: Optional[str],
@@ -309,6 +317,7 @@ def generate(
                 run_time,
                 host,
                 auth,
+                db_type,
                 dry_run,
                 custom_requirement,
                 together_api_key,
@@ -332,6 +341,7 @@ async def _async_generate(
     run_time: str,
     host: Optional[str],
     auth: bool,
+    db_type: str,
     dry_run: bool,
     custom_requirement: Optional[str],
     together_api_key: Optional[str],
@@ -343,7 +353,6 @@ async def _async_generate(
     try:
         _, api_key = _initialize_config(together_api_key)
         output_dir = _setup_output_directory(output)
-
         # Display configuration
         if ctx.obj["verbose"]:
             _display_configuration(
@@ -363,7 +372,14 @@ async def _async_generate(
         )
 
         created_files = await _generate_and_create_tests(
-            api_key, endpoints, api_info, output_dir, custom_requirement, host, auth
+            api_key,
+            endpoints,
+            api_info,
+            output_dir,
+            custom_requirement,
+            host,
+            auth,
+            db_type,
         )
 
         # Show results

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/src/devdox_ai_locust/config.py

@@ -8,7 +8,7 @@ from pydantic_settings import BaseSettings
 class Settings(BaseSettings):
     """Application settings."""
 
-    VERSION: str = "0.1.3.post1"
+    VERSION: str = "0.1.4"
 
     API_KEY: str = ""  # Fallback for backward compatibility
 

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/src/devdox_ai_locust/hybrid_loctus_generator.py

@@ -25,6 +25,7 @@ logger = logging.getLogger(__name__)
 
 
 test_data_file_path = "test_data.py"
+data_provider_path = "data_provider.py"
 
 
 @dataclass
@@ -118,6 +119,7 @@ class EnhancementProcessor:
         base_files: Dict[str, str],
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
+        db_type: str = "",
     ) -> Tuple[List[Dict[str, Any]], List[str]]:
         """Process workflow enhancements"""
         enhanced_directory_files: List[Dict[str, Any]] = []
@@ -134,10 +136,13 @@ class EnhancementProcessor:
             first_workflow = base_workflow_files[0]
             # Get the content from the dictionary - adjust key name as needed
             base_workflow_content = first_workflow.get("base_workflow.py", "")
-
         for workflow_item in directory_files:
             enhanced_workflow_item = await self._enhance_single_workflow(
-                workflow_item, base_files, base_workflow_content, grouped_endpoints
+                workflow_item,
+                base_files,
+                base_workflow_content,
+                grouped_endpoints,
+                db_type,
             )
             if enhanced_workflow_item:
                 enhanced_directory_files.append(enhanced_workflow_item["files"])
@@ -151,6 +156,7 @@ class EnhancementProcessor:
         base_files: Dict[str, str],
         base_workflow_files: str,
         grouped_endpoints: Dict[str, List[Endpoint]],
+        db_type: str = "",
     ) -> Dict[str, Any] | None:
         """Enhance a single workflow file"""
         for key, value in workflow_item.items():
@@ -164,6 +170,7 @@ class EnhancementProcessor:
                 base_workflow=base_workflow_files,
                 grouped_enpoints=workflow_endpoints_dict,
                 auth_endpoints=auth_endpoints,
+                db_type=db_type,
             )
             if enhanced_workflow:
                 return {
@@ -174,14 +181,19 @@ class EnhancementProcessor:
         return None
 
     async def process_test_data_enhancement(
-        self, base_files: Dict[str, str], endpoints: List[Endpoint]
+        self, base_files: Dict[str, str], endpoints: List[Endpoint], db_type: str = ""
     ) -> Tuple[Dict[str, str], List[str]]:
         """Process test data enhancement"""
         enhanced_files = {}
         enhancements = []
         if self.ai_config and self.ai_config.enhance_test_data:
             enhanced_test_data = await self.locust_generator.enhance_test_data_file(
-                base_files.get(test_data_file_path, ""), endpoints
+                base_files.get(test_data_file_path, ""),
+                endpoints,
+                db_type,
+                base_files.get(data_provider_path, ""),
+                base_files.get("db_config.py", ""),
+                data_provider_path,
             )
             if enhanced_test_data:
                 enhanced_files[test_data_file_path] = enhanced_test_data
@@ -289,6 +301,7 @@ class HybridLocustGenerator:
         custom_requirement: Optional[str] = None,
         target_host: Optional[str] = None,
         include_auth: bool = True,
+        db_type: str = "",
     ) -> Tuple[Dict[str, str], List[Dict[str, Any]]]:
         """
         Generate Locust tests using hybrid approach
@@ -308,6 +321,7 @@ class HybridLocustGenerator:
                     api_info,
                     include_auth=include_auth,
                     target_host=target_host,
+                    db_type=db_type,
                 )
             )
 
@@ -322,6 +336,7 @@ class HybridLocustGenerator:
                     directory_files,
                     grouped_enpoints,
                     custom_requirement,
+                    db_type,
                 )
                 if enhancement_result.success:
                     logger.info(
@@ -422,6 +437,7 @@ class HybridLocustGenerator:
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
         custom_requirement: Optional[str] = None,
+        db_type: str = "",
     ) -> EnhancementResult:
         """Enhance base files with AI - Refactored for reduced cognitive complexity"""
         start_time = asyncio.get_event_loop().time()
@@ -434,6 +450,7 @@ class HybridLocustGenerator:
                 directory_files,
                 grouped_endpoints,
                 custom_requirement,
+                db_type,
             )
 
             processing_time = asyncio.get_event_loop().time() - start_time
@@ -462,6 +479,7 @@ class HybridLocustGenerator:
         directory_files: List[Dict[str, Any]],
         grouped_endpoints: Dict[str, List[Endpoint]],
         custom_requirement: Optional[str] = None,
+        db_type: str = "",
     ) -> EnhancementResult:
         """Process all enhancements using the enhancement processor"""
         processor = EnhancementProcessor(self.ai_config, self)
@@ -476,7 +494,7 @@ class HybridLocustGenerator:
             processor.process_domain_flows_enhancement(
                 endpoints, api_info, custom_requirement
             ),
-            processor.process_test_data_enhancement(base_files, endpoints),
+            processor.process_test_data_enhancement(base_files, endpoints, db_type),
             processor.process_validation_enhancement(base_files, endpoints),
         ]
 
@@ -501,7 +519,7 @@ class HybridLocustGenerator:
                 workflow_files,
                 workflow_enhancements,
             ) = await processor.process_workflow_enhancements(
-                base_files, directory_files, grouped_endpoints
+                base_files, directory_files, grouped_endpoints, db_type
             )
             enhanced_directory_files.extend(workflow_files)
             enhancements_applied.extend(workflow_enhancements)
@@ -557,6 +575,7 @@ class HybridLocustGenerator:
         base_workflow: str,
         grouped_enpoints: Dict[str, List[Endpoint]],
         auth_endpoints: List[Endpoint],
+        db_type: str = "",
     ) -> Optional[str]:
         try:
             template = self.jinja_env.get_template("workflow.j2")
@@ -568,6 +587,7 @@ class HybridLocustGenerator:
                 base_workflow=base_workflow,
                 auth_endpoints=auth_endpoints,
                 base_content=base_content,
+                db_type=db_type,
             )
             enhanced_content = await self._call_ai_service(prompt)
             return enhanced_content
@@ -577,7 +597,13 @@ class HybridLocustGenerator:
         return ""
 
     async def enhance_test_data_file(
-        self, base_content: str, endpoints: List[Endpoint]
+        self,
+        base_content: str,
+        endpoints: List[Endpoint],
+        db_type: str = "",
+        data_provider: str = "",
+        db_config: str = "",
+        data_provider_path: str = "",
     ) -> Optional[str]:
         """Enhance test data generation with domain knowledge"""
 
@@ -592,11 +618,14 @@ class HybridLocustGenerator:
                 "base_content": base_content,
                 "schemas_info": schemas_info,
                 "endpoints": endpoints,
+                "db_type": db_type,
+                "data_provider_content": data_provider,
+                "db_config": db_config,
+                "data_provider_path": data_provider_path,
             }
 
             # Render enhanced content
             prompt = template.render(**context)
-
             enhanced_content = await self._call_ai_service(prompt)
             if enhanced_content and self._validate_python_code(enhanced_content):
                 return enhanced_content

{devdox_ai_locust-0.1.3.post1 → devdox_ai_locust-0.1.4}/src/devdox_ai_locust/locust_generator.py

@@ -12,7 +12,8 @@ import black
 from jinja2 import Environment, FileSystemLoader
 from pathlib import Path
 import logging
-from dataclasses import dataclass
+from enum import Enum
+from dataclasses import dataclass, asdict
 from datetime import datetime
 
 
@@ -21,6 +22,48 @@ from devdox_ai_locust.utils.open_ai_parser import Endpoint, Parameter
 logger = logging.getLogger(__name__)
 
 
+class DatabaseType(Enum):
+    """Supported database types for testing"""
+
+    MONGO = "mongo"
+    POSTGRES = "postgres"
+
+
+@dataclass
+class MongoDBConfig:
+    """MongoDB-specific configuration"""
+
+    use_realistic_data: str = "true"
+    enable_mongodb: str = "false"
+    use_mongodb_for_test_data: str = "false"
+    mongodb_uri: str = "mongodb://localhost:27017/"
+    mongodb_database: str = "locust_test_data"
+    MONGODB_MAX_POOL_SIZE: int = 100
+    MONGODB_MIN_POOL_SIZE: int = 10
+
+    # MongoDB Timeout Settings
+    MONGODB_CONNECT_TIMEOUT_MS: int = 5000
+    MONGODB_SERVER_SELECTION_TIMEOUT_MS: int = 5000
+    MONGODB_SOCKET_TIMEOUT_MS: int = 10000
+    MONGODB_MAX_IDLE_TIME_MS: int = 60000
+    MONGODB_WAIT_QUEUE_TIMEOUT_MS: int = 10000
+
+    # MongoDB Collection Names to be added
+
+
+@dataclass
+class PostgreSQLConfig:
+    """PostgreSQL-specific configuration"""
+
+    host: str = "localhost"
+    port: str = "5432"
+    database: str = "test_db"
+    user: str = "test_user"
+    password: str = "test_password"
+    pool_size: str = "10"
+    max_overflow: str = "20"
+
+
 @dataclass
 class TestDataConfig:
     """Configuration for test data generation"""
@@ -100,6 +143,7 @@ class LocustTestGenerator:
         api_info: Dict[str, Any],
         include_auth: bool = True,
         target_host: Optional[str] = None,
+        db_type: str = "",
     ) -> Tuple[Dict[str, str], List[Dict[str, Any]], Dict[str, List[Endpoint]]]:
         """
         Generate complete Locust test suite from parsed endpoints
@@ -121,14 +165,23 @@ class LocustTestGenerator:
                 "locustfile.py": self._generate_main_locustfile(
                     endpoints, api_info, list(grouped_enpoint.keys())
                 ),
-                "test_data.py": self._generate_test_data_file(),
+                "test_data.py": self._generate_test_data_file(db_type),
                 "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),
+                "README.md": self._generate_readme_file(api_info, db_type),
+                ".env.example": self._generate_env_example(
+                    api_info, target_host, db_type
+                ),
             }
+            if db_type != "":
+                self.generated_files["db_config.py"] = self._generate_db_file(
+                    db_type, "db_config.py.j2"
+                )
+                self.generated_files["data_provider.py"] = self._generate_db_file(
+                    db_type, "data_provider.py.j2"
+                )
 
             return self.generated_files, workflows_files, grouped_enpoint
         except Exception as e:
@@ -639,10 +692,13 @@ class LocustTestGenerator:
 
     '''
 
-    def _generate_test_data_file(self) -> str:
+    def _generate_test_data_file(self, db_type: str = "") -> str:
         """Generate test_data.py file content"""
+        data_provider_content = None
+        if db_type == DatabaseType.MONGO.value:
+            data_provider_content = "mongo_data_provider"
         template = self.jinja_env.get_template("test_data.py.j2")
-        return template.render()
+        return template.render(data_provider_content=data_provider_content)
 
     def _generate_config_file(self, api_info: Dict[str, Any]) -> str:
         """Generate config.py file content"""
@@ -665,8 +721,16 @@ class LocustTestGenerator:
         content = template.render()
         return content
 
+    def _generate_db_file(self, db_type: str, file_name: str) -> str:
+        """Generate db file content"""
+        template = self.jinja_env.get_template(db_type + "/" + file_name)
+        return template.render()
+
     def _generate_env_example(
-        self, api_info: Dict[str, Any], target_host: Optional[str] = None
+        self,
+        api_info: Dict[str, Any],
+        target_host: Optional[str] = None,
+        db_type: str = "",
     ) -> str:
         """Generate .env.example file content"""
         try:
@@ -686,12 +750,15 @@ class LocustTestGenerator:
                 "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",
             }
 
+            if db_type == DatabaseType.MONGO.value:
+                mongodb_config = MongoDBConfig()
+                environment_vars.update(asdict(mongodb_config))
+
             context = {
                 "environment_vars": environment_vars,
                 "api_info": api_info,
@@ -706,15 +773,23 @@ class LocustTestGenerator:
             logger.error(f"❌ Failed to generate .env.example from template: {e}")
             return ""
 
-    def _generate_readme_file(self, api_info: Dict[str, Any]) -> str:
+    def _generate_readme_file(self, api_info: Dict[str, Any], db_type: str = "") -> str:
         try:
             # Get the template
             template = self.jinja_env.get_template("readme.md.j2")
+            db_using = ""
+            if db_type == DatabaseType.MONGO.value:
+                template_db = self.jinja_env.get_template(
+                    DatabaseType.MONGO.value + "/db_integration.j2"
+                )
+
+                db_using = template_db.render()
 
             # Prepare template context
             context = {
                 "api_info": api_info,
                 "generated_at": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
+                "db_using": db_using,
             }
 
             # Render the template