devdox-ai-locust 0.1.3.post1__py3-none-any.whl → 0.1.4__py3-none-any.whl

devdox_ai_locust/templates/mongo/data_provider.py.j2

@@ -0,0 +1,303 @@
+"""
+MongoDB Data Provider for Locust Load Testing
+
+Provides efficient data retrieval from MongoDB with:
+- Random sampling for realistic load distribution
+- Caching to reduce database queries
+- Fallback to generated data if MongoDB unavailable
+- Smart data rotation to prevent hotspots
+"""
+
+import logging
+import random
+from typing import Dict, List, Any, Optional
+from datetime import datetime
+from collections import deque
+import threading
+import uuid
+from db_config import mongo_config
+
+logger = logging.getLogger(__name__)
+
+
+class MongoDataProvider:
+    """
+    Provides test data from MongoDB with intelligent caching and fallback.
+
+    This class is designed for high-concurrency load testing:
+    - Thread-safe caching mechanism
+    - Rotating cache to prevent data hotspots
+    - Automatic fallback to generated data
+    - Periodic cache refresh
+    """
+
+    def __init__(self, cache_size: int = 1000, cache_ttl_seconds: int = 300):
+        """
+        Initialize MongoDB data provider
+
+        Args:
+            cache_size: Maximum items to cache per collection
+            cache_ttl_seconds: Cache time-to-live in seconds
+        """
+        self.cache_size = cache_size
+        self.cache_ttl_seconds = cache_ttl_seconds
+
+        # Thread-safe caches with rotation
+        self._cache = {}
+        self._cache_timestamps = {}
+        self._cache_lock = threading.Lock()
+
+        # Usage tracking for intelligent preloading
+        self._access_counts = {}
+
+        # Statistics
+        self.stats = {
+            "cache_hits": 0,
+            "cache_misses": 0,
+            "db_queries": 0,
+            "fallback_generations": 0,
+        }
+
+    def _get_cache_key(self, collection: str, query: Dict) -> str:
+        """Generate cache key from collection and query"""
+        query_str = str(sorted(query.items()))
+        return f"{collection}:{query_str}"
+
+    def _is_cache_valid(self, cache_key: str) -> bool:
+        """Check if cache is still valid"""
+        if cache_key not in self._cache_timestamps:
+            return False
+
+        age = (datetime.now() - self._cache_timestamps[cache_key]).total_seconds()
+        return age < self.cache_ttl_seconds
+
+    def _update_cache(self, cache_key: str, data: List[Dict]):
+        """Update cache with new data"""
+        with self._cache_lock:
+            # Limit cache size using deque for efficient rotation
+            if len(data) > self.cache_size:
+                self._cache[cache_key] = deque(
+                    random.sample(data, self.cache_size),
+                    maxlen=self.cache_size,
+                )
+            else:
+                self._cache[cache_key] = deque(data, maxlen=self.cache_size)
+
+            self._cache_timestamps[cache_key] = datetime.now()
+
+    def get_random_document(
+        self,
+        collection_name: str,
+        query: Dict = None,
+        projection: Dict = None,
+    ) -> Optional[Dict]:
+        """
+        Get a random document from MongoDB collection with caching.
+
+        Args:
+            collection_name: Name of the collection (e.g., 'pets', 'users')
+            query: MongoDB query filter
+            projection: Fields to include/exclude
+
+        Returns:
+            Random document from collection, or None if not found
+        """
+        if not mongo_config.enable_mongodb:
+            logger.debug("MongoDB disabled, using generated data")
+            return self._generate_fallback_data(collection_name)
+
+        query = query or {}
+        cache_key = self._get_cache_key(collection_name, query)
+
+        # Try cache first
+        if cache_key in self._cache and self._is_cache_valid(cache_key):
+            self.stats["cache_hits"] += 1
+            with self._cache_lock:
+                cached_data = list(self._cache[cache_key])
+                if cached_data:
+                    # Rotate the deque for better distribution
+                    item = cached_data[0]
+                    self._cache[cache_key].rotate(-1)
+                    return item
+
+        # Cache miss - query database
+        self.stats["cache_misses"] += 1
+        try:
+            collection = mongo_config.get_collection(collection_name)
+            if not collection:
+                logger.warning(f"Collection {collection_name} not available")
+                return self._generate_fallback_data(collection_name)
+
+            self.stats["db_queries"] += 1
+
+            # Fetch multiple documents for caching
+            documents = list(collection.find(query, projection).limit(self.cache_size))
+
+            if documents:
+                # Update cache
+                self._update_cache(cache_key, documents)
+
+                # Return random document
+                return random.choice(documents)
+            else:
+                logger.warning(f"No documents found in {collection_name} with query {query}")
+                return self._generate_fallback_data(collection_name)
+
+        except Exception as e:
+            logger.error(f"Error fetching from MongoDB {collection_name}: {e}")
+            self.stats["fallback_generations"] += 1
+            return self._generate_fallback_data(collection_name)
+
+    def get_document(
+        self,
+        collection_name: str,
+        query: Dict = None,
+        projection: Dict = None,
+        sort: List[tuple] = None,
+    ) -> Optional[Dict]:
+        """
+        Get a single specific document from MongoDB (not random).
+        Returns the first matching document based on query and sort order.
+        """
+        if not mongo_config.enable_mongodb:
+            logger.debug("MongoDB disabled, using generated data")
+            return self._generate_fallback_data(collection_name)
+
+        query = query or {}
+
+        try:
+            collection = mongo_config.get_collection(collection_name)
+            if not collection:
+                logger.warning(f"Collection {collection_name} not available")
+                return self._generate_fallback_data(collection_name)
+
+            self.stats["db_queries"] += 1
+
+            # Build the query
+            cursor = collection.find(query, projection)
+
+            # Apply sort if specified
+            if sort:
+                cursor = cursor.sort(sort)
+
+            # Get first document
+            document = cursor.limit(1).next() if cursor else None
+
+            if document:
+                return document
+            else:
+                logger.warning(f"No document found in {collection_name} with query {query}")
+                return self._generate_fallback_data(collection_name)
+
+        except StopIteration:
+            logger.warning(f"No document found in {collection_name} with query {query}")
+            return self._generate_fallback_data(collection_name)
+        except Exception as e:
+            logger.error(f"Error fetching from MongoDB {collection_name}: {e}")
+            self.stats["fallback_generations"] += 1
+            return self._generate_fallback_data(collection_name)
+
+    def get_multiple_documents(
+        self,
+        collection_name: str,
+        count: int = 10,
+        query: Dict = None,
+        projection: Dict = None,
+        random: bool = True,
+    ) -> List[Dict]:
+        """
+        Get multiple random documents from collection.
+        """
+        collection = mongo_config.get_collection(collection_name)
+        if random:
+            # ✅ SINGLE EFFICIENT QUERY using MongoDB's $sample
+            pipeline = []
+            if query:
+                pipeline.append({"$match": query})
+            pipeline.append({"$sample": {"size": count}})
+            if projection:
+                pipeline.append({"$project": projection})
+
+            documents = list(collection.aggregate(pipeline))
+            return documents
+        else:
+            try:
+                collection = mongo_config.get_collection(collection_name)
+                if not collection:
+                    logger.warning(f"Collection {collection_name} not available")
+                    return [self._generate_fallback_data(collection_name) for _ in range(count)]
+
+                query = query or {}
+                self.stats["db_queries"] += 1
+
+                documents = list(collection.find(query, projection).limit(count))
+
+                if documents:
+                    return documents
+                else:
+                    logger.warning(f"No documents found in {collection_name} with query {query}")
+                    return [self._generate_fallback_data(collection_name) for _ in range(count)]
+
+            except Exception as e:
+                logger.error(f"Error fetching from MongoDB {collection_name}: {e}")
+                return [self._generate_fallback_data(collection_name) for _ in range(count)]
+
+    def _generate_fallback_data(self, collection_name: str) -> Dict:
+        """
+        Generate fallback data when MongoDB is unavailable.
+        """
+        self.stats["fallback_generations"] += 1
+
+        # Generic fallback
+        return {
+            "_id":  str(uuid.uuid4()),
+            "data": f"fallback_{collection_name}",
+        }
+
+    def clear_cache(self):
+        """Clear all cached data"""
+        with self._cache_lock:
+            self._cache.clear()
+            self._cache_timestamps.clear()
+        logger.info("Cache cleared")
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get usage statistics"""
+        total_requests = self.stats["cache_hits"] + self.stats["cache_misses"]
+
+        return {
+            **self.stats,
+            "cache_hit_rate": (
+                self.stats["cache_hits"] / total_requests * 100 if total_requests > 0 else 0
+            ),
+            "cache_size": sum(len(cache) for cache in self._cache.values()),
+            "cached_collections": len(self._cache),
+        }
+
+    def preload_cache(self, collection_name: str, query: Dict = None):
+        """
+        Preload cache for a collection to improve initial performance.
+        """
+        logger.info(f"Preloading cache for {collection_name}...")
+        try:
+            collection = mongo_config.get_collection(collection_name)
+            if not collection:
+                logger.warning(f"Collection {collection_name} not available for preload")
+                return
+
+            query = query or {}
+            documents = list(collection.find(query).limit(self.cache_size))
+
+            if documents:
+                cache_key = self._get_cache_key(collection_name, query)
+                self._update_cache(cache_key, documents)
+                logger.info(f"✅ Preloaded {len(documents)} documents for {collection_name}")
+            else:
+                logger.warning(f"No documents found to preload for {collection_name}")
+
+        except Exception as e:
+            logger.error(f"Failed to preload cache for {collection_name}: {e}")
+
+
+# Global instance
+mongo_data_provider = MongoDataProvider()

devdox_ai_locust/templates/mongo/db_config.py.j2

@@ -0,0 +1,271 @@
+"""
+MongoDB Configuration for Load Testing
+
+This module provides MongoDB connection management with:
+- Connection pooling optimized for concurrent load testing
+- Singleton pattern to prevent connection exhaustion
+- Proper error handling and retry logic
+- Environment-based configuration
+"""
+
+import os
+import logging
+from typing import Optional, Dict, Any, List
+from pymongo import MongoClient
+from pymongo.errors import ConnectionFailure, ServerSelectionTimeoutError
+from dotenv import load_dotenv
+import time
+
+load_dotenv()
+logger = logging.getLogger(__name__)
+
+
+class MongoDBConfig:
+    """MongoDB configuration with connection pooling for load testing"""
+
+    _instance = None
+    _client = None
+    _database = None
+
+    def __new__(cls):
+        """Singleton pattern to ensure single connection pool"""
+        if cls._instance is None:
+            cls._instance = super(MongoDBConfig, cls).__new__(cls)
+            cls._instance._initialize()
+        return cls._instance
+
+    def _initialize(self):
+        """Initialize MongoDB connection with proper pooling"""
+        # MongoDB Configuration
+        self.mongo_uri = os.getenv("MONGODB_URI", "mongodb://localhost:27017/")
+        self.database_name = os.getenv("MONGODB_DATABASE", "locust_test_data")
+        
+        # Connection Pool Settings - CRITICAL FOR LOAD TESTING
+        # These values are optimized for concurrent Locust users
+        self.max_pool_size = int(os.getenv("MONGODB_MAX_POOL_SIZE", "100"))
+        self.min_pool_size = int(os.getenv("MONGODB_MIN_POOL_SIZE", "10"))
+        
+        # Timeout Settings
+        self.connect_timeout_ms = int(os.getenv("MONGODB_CONNECT_TIMEOUT_MS", "5000"))
+        self.server_selection_timeout_ms = int(
+            os.getenv("MONGODB_SERVER_SELECTION_TIMEOUT_MS", "5000")
+        )
+        self.socket_timeout_ms = int(os.getenv("MONGODB_SOCKET_TIMEOUT_MS", "10000"))
+        
+        # Wait Queue Settings - Prevents connection exhaustion
+        self.max_idle_time_ms = int(os.getenv("MONGODB_MAX_IDLE_TIME_MS", "60000"))
+        self.wait_queue_timeout_ms = int(
+            os.getenv("MONGODB_WAIT_QUEUE_TIMEOUT_MS", "10000")
+        )
+        
+        # Feature Flags
+        self.enable_mongodb = os.getenv("ENABLE_MONGODB", "false").lower() == "true"
+        self.use_mongodb_for_test_data = (
+            os.getenv("USE_MONGODB_FOR_TEST_DATA", "false").lower() == "true"
+        )
+        
+        # Collections
+        self.collections = {
+            "pets": os.getenv("MONGODB_COLLECTION_PETS", "pets"),
+            "users": os.getenv("MONGODB_COLLECTION_USERS", "users"),
+            "orders": os.getenv("MONGODB_COLLECTION_ORDERS", "orders"),
+            "test_sessions": os.getenv("MONGODB_COLLECTION_SESSIONS", "test_sessions"),
+        }
+
+    def get_client(self) -> Optional[MongoClient]:
+        """
+        Get MongoDB client with connection pooling.
+        
+        CRITICAL: Returns same client instance for all Locust workers
+        to share connection pool efficiently.
+        """
+        if not self.enable_mongodb:
+            logger.warning("MongoDB is disabled. Enable with ENABLE_MONGODB=true")
+            return None
+
+        if self._client is None:
+            try:
+                logger.info("Initializing MongoDB connection pool...")
+                logger.info(f"Pool size: {self.min_pool_size} to {self.max_pool_size}")
+                
+                self._client = MongoClient(
+                    self.mongo_uri,
+                    # Connection Pool Configuration
+                    maxPoolSize=self.max_pool_size,
+                    minPoolSize=self.min_pool_size,
+                    maxIdleTimeMS=self.max_idle_time_ms,
+                    waitQueueTimeoutMS=self.wait_queue_timeout_ms,
+                    
+                    # Timeout Configuration
+                    connectTimeoutMS=self.connect_timeout_ms,
+                    serverSelectionTimeoutMS=self.server_selection_timeout_ms,
+                    socketTimeoutMS=self.socket_timeout_ms,
+                    
+                    # Resilience Settings
+                    retryWrites=True,
+                    retryReads=True,
+                    
+                    # Application Name for monitoring
+                    appname="locust-load-test",
+                )
+                
+                # Verify connection
+                self._client.admin.command("ping")
+                logger.info(f"✅ MongoDB connected successfully to {self.database_name}")
+                
+            except (ConnectionFailure, ServerSelectionTimeoutError) as e:
+                logger.error(f"❌ Failed to connect to MongoDB: {e}")
+                logger.error("Please verify MongoDB is running and connection string is correct")
+                self._client = None
+                raise
+            except Exception as e:
+                logger.error(f"❌ Unexpected MongoDB connection error: {e}")
+                self._client = None
+                raise
+
+        return self._client
+
+    def get_database(self):
+        """Get database instance"""
+        if self._database is None:
+            client = self.get_client()
+            if client:
+                self._database = client[self.database_name]
+        return self._database
+
+    def get_collection(self, collection_name: str):
+        """Get collection instance"""
+        db = self.get_database()
+        if db:
+            # Use the mapped collection name if it exists
+            actual_name = self.collections.get(collection_name, collection_name)
+            return db[actual_name]
+        return None
+
+    def close(self):
+        """Close MongoDB connection"""
+        if self._client:
+            logger.info("Closing MongoDB connection...")
+            self._client.close()
+            self._client = None
+            self._database = None
+            logger.info("✅ MongoDB connection closed")
+
+    def validate_config(self) -> List[str]:
+        """
+        Validate MongoDB configuration and return warnings/errors
+        
+        Returns:
+            List of configuration issues
+        """
+        issues = []
+
+        # Pool size validation
+        if self.max_pool_size < self.min_pool_size:
+            issues.append(
+                f"⚠️  max_pool_size ({self.max_pool_size}) is less than "
+                f"min_pool_size ({self.min_pool_size})"
+            )
+
+        # Check if pool size is adequate for Locust users
+        # Rule of thumb: max_pool_size should be >= number of Locust users
+        locust_users = int(os.getenv("LOCUST_USERS", "50"))
+        if self.max_pool_size < locust_users:
+            issues.append(
+                f"⚠️  CRITICAL: max_pool_size ({self.max_pool_size}) is less than "
+                f"LOCUST_USERS ({locust_users}). This WILL cause connection "
+                f"exhaustion and test failures. Recommended: {locust_users * 2}"
+            )
+
+        # Timeout validation
+        if self.connect_timeout_ms < 1000:
+            issues.append(
+                f"⚠️  connect_timeout_ms ({self.connect_timeout_ms}) is very low. "
+                "May cause connection failures under load."
+            )
+
+        if self.wait_queue_timeout_ms < 5000:
+            issues.append(
+                f"⚠️  wait_queue_timeout_ms ({self.wait_queue_timeout_ms}) is low. "
+                "Users may get connection timeouts under heavy load."
+            )
+
+        return issues
+
+    def get_connection_stats(self) -> Dict[str, Any]:
+        """Get current connection pool statistics"""
+        client = self.get_client()
+        if not client:
+            return {"status": "disabled"}
+
+        try:
+            # Get server status for connection pool info
+            server_status = client.admin.command("serverStatus")
+            connections = server_status.get("connections", {})
+            
+            return {
+                "status": "connected",
+                "current_connections": connections.get("current", 0),
+                "available_connections": connections.get("available", 0),
+                "total_created": connections.get("totalCreated", 0),
+                "active": connections.get("active", 0),
+                "configured_max": self.max_pool_size,
+                "configured_min": self.min_pool_size,
+            }
+        except Exception as e:
+            logger.error(f"Failed to get connection stats: {e}")
+            return {"status": "error", "error": str(e)}
+
+
+# Global singleton instance
+mongo_config = MongoDBConfig()
+
+
+def test_connection():
+    """Test MongoDB connection and display configuration"""
+    print("\n" + "=" * 60)
+    print("MongoDB Connection Test")
+    print("=" * 60)
+    
+    # Display configuration
+    print(f"\nConnection URI: {mongo_config.mongo_uri}")
+    print(f"Database: {mongo_config.database_name}")
+    print(f"Pool Size: {mongo_config.min_pool_size} - {mongo_config.max_pool_size}")
+    print(f"Connect Timeout: {mongo_config.connect_timeout_ms}ms")
+    print(f"Socket Timeout: {mongo_config.socket_timeout_ms}ms")
+    
+    # Check for configuration issues
+    issues = mongo_config.validate_config()
+    if issues:
+        print("\n⚠️  Configuration Issues:")
+        for issue in issues:
+            print(f"  {issue}")
+    else:
+        print("\n✅ Configuration looks good!")
+    
+    # Try to connect
+    print("\nAttempting to connect...")
+    try:
+        client = mongo_config.get_client()
+        if client:
+            stats = mongo_config.get_connection_stats()
+            print("\n✅ Connection successful!")
+            print(f"\nConnection Stats:")
+            for key, value in stats.items():
+                print(f"  {key}: {value}")
+            
+            # List collections
+            db = mongo_config.get_database()
+            collections = db.list_collection_names()
+            print(f"\nAvailable collections: {', '.join(collections) if collections else 'None'}")
+        else:
+            print("\n❌ Connection failed - MongoDB is disabled")
+            
+    except Exception as e:
+        print(f"\n❌ Connection failed: {e}")
+    
+    print("\n" + "=" * 60)
+
+
+if __name__ == "__main__":
+    test_connection()