run_performance_tests.py

#!/usr/bin/env python3
"""
Performance test script for Flexo API
Supports sequential and concurrent test execution with performance metrics and charts
"""

import json
import subprocess
import time
import argparse
import asyncio
import aiohttp
from pathlib import Path
from datetime import datetime
from typing import List, Dict, Any
from collections import defaultdict
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from matplotlib.figure import Figure
from dotenv import load_dotenv
import os

# Configuration
ENDPOINT = "URL"
TEST_MODEL_FILE = "./sysml/Flashlight.json"

# Default timeout (in seconds) for all HTTP requests (overridable via CLI)
DEFAULT_REQUEST_TIMEOUT_SECONDS = 900  # 15 minutes


class PerformanceTracker:
    """Tracks performance metrics and errors"""
    
    def __init__(self):
        self.metrics = {
            'create_project': [],
            'commit': [],
            'get_elements': [],
            'get_projects': []
        }
        self.errors = []  # Track all errors
        self.timestamps = []
        self.start_time = None
        
    def record_metric(self, operation: str, duration_ms: float, metadata: Dict = None):
        """Record a performance metric"""
        if operation in self.metrics:
            self.metrics[operation].append({
                'duration_ms': duration_ms,
                'timestamp': time.time(),
                'metadata': metadata or {}
            })
    
    def record_error(self, operation: str, iteration: int, error_type: str, error_message: str, metadata: Dict = None):
        """Record an error"""
        self.errors.append({
            'operation': operation,
            'iteration': iteration,
            'error_type': error_type,
            'error_message': error_message,
            'timestamp': time.time(),
            'metadata': metadata or {}
        })
    
    def get_statistics(self) -> Dict[str, Dict]:
        """Calculate statistics for each operation"""
        stats = {}
        for operation, measurements in self.metrics.items():
            if not measurements:
                continue
            durations = [m['duration_ms'] for m in measurements]
            stats[operation] = {
                'count': len(durations),
                'min': min(durations),
                'max': max(durations),
                'avg': sum(durations) / len(durations),
                'total': sum(durations)
            }
        return stats
    
    def get_error_statistics(self) -> Dict[str, Any]:
        """Get error statistics"""
        if not self.errors:
            return {'total_errors': 0, 'errors_by_operation': {}, 'errors_by_type': {}}
        
        errors_by_operation = defaultdict(int)
        errors_by_type = defaultdict(int)
        
        for error in self.errors:
            errors_by_operation[error['operation']] += 1
            errors_by_type[error['error_type']] += 1
        
        return {
            'total_errors': len(self.errors),
            'errors_by_operation': dict(errors_by_operation),
            'errors_by_type': dict(errors_by_type),
            'errors': self.errors
        }


class FlexoTestRunner:
    """Runs Flexo API performance tests"""
    
    def __init__(self, endpoint: str, bearer_token: str, test_model_path: str,
                 request_timeout: int = DEFAULT_REQUEST_TIMEOUT_SECONDS):
        self.endpoint = endpoint
        self.bearer_token = bearer_token
        self.test_model_path = test_model_path
        self.request_timeout = request_timeout
        self.tracker = PerformanceTracker()
        self.project_ids = []
        self.commit_ids = []  # Track commit IDs for getting elements
        self.element_counts = []  # Track element counts per iteration for triple estimation
    
    def _client_timeout(self) -> aiohttp.ClientTimeout:
        """Return a configured aiohttp timeout"""
        return aiohttp.ClientTimeout(total=self.request_timeout, sock_read=self.request_timeout)
        
    def load_test_model(self) -> str:
        """Load test model JSON"""
        with open(self.test_model_path, 'r') as f:
            return f.read()
    
    async def create_project(self, session: aiohttp.ClientSession, project_name: str) -> str:
        """Create a new project"""
        url = f"{self.endpoint}/projects"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.bearer_token}"
        }
        payload = {
            "name": project_name,
            "description": f"Performance test project - {project_name}"
        }
        
        start_time = time.time()
        async with session.post(
            url,
            json=payload,
            headers=headers,
            timeout=self._client_timeout(),
        ) as response:
            duration_ms = (time.time() - start_time) * 1000
            
            # Try to parse JSON, but handle non-JSON responses
            try:
                data = await response.json()
            except aiohttp.client_exceptions.ContentTypeError as e:
                try:
                    error_text = await response.text()
                    error_preview = error_text[:200] if error_text else "No error message"
                except:
                    error_preview = f"Could not read response: {str(e)}"
                error_msg = f"Failed to create project: {response.status} - Server returned non-JSON response: {error_preview}"
                self.tracker.record_error('create_project', 0, 'NonJsonResponse', error_msg, {'status': response.status})
                raise Exception(error_msg)
            except Exception as e:
                try:
                    error_text = await response.text()
                    error_preview = error_text[:200] if error_text else "No error message"
                except:
                    error_preview = f"Could not read response: {str(e)}"
                error_msg = f"Failed to create project: {response.status} - JSON parse error: {str(e)} - Response: {error_preview}"
                self.tracker.record_error('create_project', 0, 'JsonParseError', error_msg, {'status': response.status})
                raise Exception(error_msg)
            
            # Extract project ID - prioritize @id as shown in user's example response
            project_id = data.get('@id') or data.get('id') or data.get('projectId')
            
            # Debug: Log the full response if project_id is missing
            if not project_id:
                print(f"    ⚠️  DEBUG: Could not extract project_id from response. Response keys: {list(data.keys()) if isinstance(data, dict) else 'not a dict'}")
                print(f"    ⚠️  DEBUG: Full response: {data}")
            
            self.tracker.record_metric('create_project', duration_ms, {
                'project_id': project_id,
                'status': response.status
            })
            
            if response.status in (200, 201) and project_id:
                # Don't add to self.project_ids here - let the caller do it with proper locking
                # This prevents race conditions in concurrent mode
                return project_id
            else:
                error_msg = f"Failed to create project: {response.status} - {data}"
                if not project_id:
                    error_msg += f" (Could not extract project_id from response)"
                self.tracker.record_error('create_project', 0, 'CreateProjectFailed', error_msg, {'status': response.status, 'response_data': str(data)[:200]})
                raise Exception(error_msg)
    
    async def commit_model(self, session: aiohttp.ClientSession, project_id: str, model_json: str) -> str:
        """Commit model to a project and return commit ID"""
        # Validate project_id is not None or empty
        if not project_id:
            raise Exception("Cannot commit: project_id is None or empty")
        
        # Validate project_id format
        if not isinstance(project_id, str) or len(project_id) < 30:
            raise Exception(f"Cannot commit: Invalid project_id format: {project_id} (type: {type(project_id)}, length: {len(project_id) if project_id else 0})")
        
        url = f"{self.endpoint}/projects/{project_id}/commits"
        headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.bearer_token}"
        }
        
        # Parse and wrap in Commit structure
        model_data = json.loads(model_json)
        commit_payload = {
            "@type": "Commit",
            "description": "Performance test commit",
            "change": model_data
        }
        
        start_time = time.time()
        try:
            # Use a long timeout (10 minutes) to allow very long-running commits
            async with session.post(
                url,
                json=commit_payload,
                headers=headers,
                timeout=self._client_timeout(),
            ) as response:
                duration_ms = (time.time() - start_time) * 1000
                
                # Try to parse JSON, but handle non-JSON responses (e.g., HTML error pages)
                try:
                    data = await response.json()
                except aiohttp.client_exceptions.ContentTypeError as e:
                    # Response is not JSON (likely HTML error page from timeout/error)
                    try:
                        error_text = await response.text()
                        error_preview = error_text[:200] if error_text else "No error message"
                    except:
                        error_preview = f"Could not read response: {str(e)}"
                    
                    self.tracker.record_metric('commit', duration_ms, {
                        'project_id': project_id,
                        'commit_id': None,
                        'status': response.status,
                        'error': f'non_json_response: {error_preview}'
                    })
                    error_msg = f"Failed to commit: {response.status} - Server returned non-JSON response (likely timeout/error page): {error_preview}"
                    self.tracker.record_error('commit', 0, 'NonJsonResponse', error_msg, {'project_id': project_id, 'status': response.status})
                    raise Exception(error_msg)
                except Exception as e:
                    # Other JSON parsing errors
                    try:
                        error_text = await response.text()
                        error_preview = error_text[:200] if error_text else "No error message"
                    except:
                        error_preview = f"Could not read response: {str(e)}"
                    
                    self.tracker.record_metric('commit', duration_ms, {
                        'project_id': project_id,
                        'commit_id': None,
                        'status': response.status,
                        'error': f'json_parse_error: {str(e)}'
                    })
                    error_msg = f"Failed to commit: {response.status} - JSON parse error: {str(e)} - Response: {error_preview}"
                    self.tracker.record_error('commit', 0, 'JsonParseError', error_msg, {'project_id': project_id, 'status': response.status})
                    raise Exception(error_msg)
                
                commit_id = data.get('@id') or data.get('id') or data.get('commitId')
                self.tracker.record_metric('commit', duration_ms, {
                    'project_id': project_id,
                    'commit_id': commit_id,
                    'status': response.status
                })
                
                # Handle 404 specifically - project might have been deleted
                if response.status == 404:
                    error_msg = f"Failed to commit: Project {project_id} does not exist (404) - project may have been deleted"
                    self.tracker.record_error('commit', 0, 'ProjectNotFound', error_msg, {'project_id': project_id, 'status': 404})
                    raise Exception(error_msg)
                
                if response.status not in (200, 201):
                    error_msg = f"Failed to commit: {response.status} - {data}"
                    self.tracker.record_error('commit', 0, 'CommitFailed', error_msg, {'project_id': project_id, 'status': response.status})
                    raise Exception(error_msg)
                
                if not commit_id:
                    error_msg = f"Commit succeeded but no commit ID found in response: {data}"
                    self.tracker.record_error('commit', 0, 'MissingCommitId', error_msg, {'project_id': project_id})
                    raise Exception(error_msg)
                
                return commit_id
        except asyncio.TimeoutError:
            # Timeout occurred - the request might still be processing on the server
            # IMPORTANT: The server may continue processing the commit even though the client timed out
            # We should NOT delete this project immediately - wait for server to finish
            duration_ms = (time.time() - start_time) * 1000
            error_msg = f"Timeout while committing to project {project_id} (request may still be processing on server - project will not be deleted immediately)"
            self.tracker.record_metric('commit', duration_ms, {
                'project_id': project_id,
                'commit_id': None,
                'status': 'timeout',
                'error': 'timeout'
            })
            self.tracker.record_error('commit', 0, 'TimeoutError', error_msg, {'project_id': project_id})
            # Mark this project as "pending" - don't delete it in cleanup
            # The project will be cleaned up in a future run or manually
            raise Exception(error_msg)
        except Exception as e:
            # Re-raise other exceptions
            raise
    
    async def get_elements(self, session: aiohttp.ClientSession, project_id: str, commit_id: str = None, max_retries: int = 20, retry_delay: float = 3.0, iteration_prefix: str = ""):
        """Get elements from a project commit with retry logic (commit graph may need time to build)"""
        if commit_id:
            url = f"{self.endpoint}/projects/{project_id}/commits/{commit_id}/elements"
        else:
            # Fallback to old endpoint if commit_id not provided
            url = f"{self.endpoint}/projects/{project_id}/elements"
        headers = {
            "Authorization": f"Bearer {self.bearer_token}"
        }
        
        prefix = f"{iteration_prefix} " if iteration_prefix else ""
        
        # Retry loop: commit graph may need time to build
        for attempt in range(max_retries):
            if attempt > 0:
                await asyncio.sleep(retry_delay)
                print(f"{prefix}  🔄 Retry attempt {attempt + 1}/{max_retries} (waiting for commit graph to be ready)...")
            
            start_time = time.time()
            try:
                async with session.get(
                    url,
                    headers=headers,
                    timeout=self._client_timeout(),
                ) as response:
                    duration_ms = (time.time() - start_time) * 1000
                    
                    # Handle different response statuses
                    if response.status == 404:
                        # Commit graph may still be building - retry if we have attempts left
                        if attempt < max_retries - 1:
                            # Will retry on next iteration
                            continue
                        else:
                            # Last attempt failed
                            try:
                                error_text = await response.text()
                            except:
                                error_text = "No error message"
                            print(f"{prefix}  ⚠️  Error: 404 after {max_retries} attempts - {error_text[:200]}")
                            # Extract iteration from prefix if available (format: "[Iter N] ")
                            iteration = 0
                            if iteration_prefix:
                                try:
                                    # Extract number from "[Iter N]"
                                    import re
                                    match = re.search(r'\[Iter (\d+)\]', iteration_prefix)
                                    if match:
                                        iteration = int(match.group(1))
                                except:
                                    pass
                            self.tracker.record_error('get_elements', iteration, 'EndpointNotFound', f'404 after {max_retries} attempts: {error_text[:200]}', {'project_id': project_id, 'commit_id': commit_id})
                            self.tracker.record_metric('get_elements', duration_ms, {
                                'project_id': project_id,
                                'commit_id': commit_id,
                                'element_count': 0,
                                'status': response.status,
                                'error': f'endpoint_not_found_after_retries: {error_text[:200]}'
                            })
                            return 0
                
                    if response.status != 200:
                        # Try to get error message
                        try:
                            error_text = await response.text()
                        except:
                            error_text = f"Status {response.status}"
                        # If not 200 and not 404 (already handled), return error
                        # Extract iteration from prefix if available
                        iteration = 0
                        if iteration_prefix:
                            try:
                                import re
                                match = re.search(r'\[Iter (\d+)\]', iteration_prefix)
                                if match:
                                    iteration = int(match.group(1))
                            except:
                                pass
                        self.tracker.record_error('get_elements', iteration, 'GetElementsFailed', f'Status {response.status}: {error_text[:200]}', {'project_id': project_id, 'status': response.status})
                        self.tracker.record_metric('get_elements', duration_ms, {
                            'project_id': project_id,
                            'element_count': 0,
                            'status': response.status,
                            'error': error_text[:200]
                        })
                        print(f"{prefix}  ⚠️  Warning: Get elements returned {response.status}: {error_text[:100]}")
                        return 0  # Return gracefully instead of raising
                    
                    # Success! Parse JSON response
                    try:
                        data = await response.json()
                        element_count = len(data) if isinstance(data, list) else len(data.get('elements', []))
                        
                        if attempt > 0:
                            print(f"{prefix}  ✓ Elements available after {attempt + 1} attempt(s)")
                        
                        self.tracker.record_metric('get_elements', duration_ms, {
                            'project_id': project_id,
                            'commit_id': commit_id,
                            'element_count': element_count,
                            'status': response.status,
                            'retry_count': attempt
                        })
                        
                        return element_count
                    except Exception as e:
                        # If JSON parsing fails, try text
                        print(f"{prefix}  ⚠️  Debug: JSON parse error: {type(e).__name__}: {str(e)}")
                        try:
                            text = await response.text()
                            print(f"{prefix}  🔍 Debug: Response text length: {len(text)} chars")
                            print(f"{prefix}  🔍 Debug: Response text (first 500 chars): {text[:500]}")
                        except Exception as text_error:
                            print(f"{prefix}  ⚠️  Debug: Could not read response text: {str(text_error)}")
                        self.tracker.record_metric('get_elements', duration_ms, {
                            'project_id': project_id,
                            'commit_id': commit_id,
                            'element_count': 0,
                            'status': response.status,
                            'error': f'json_parse_error: {str(e)}'
                        })
                        print(f"{prefix}  ⚠️  Warning: Could not parse JSON response: {str(e)}")
                        return 0
            except asyncio.TimeoutError:
                if attempt < max_retries - 1:
                    print(f"{prefix}  ⚠️  Request timed out, retrying...")
                    continue
                else:
                    print(f"{prefix}  ⚠️  Error: Request timed out after {max_retries} attempts")
                    return 0
            except Exception as e:
                if attempt < max_retries - 1:
                    print(f"{prefix}  ⚠️  Exception (will retry): {type(e).__name__}: {str(e)}")
                    continue
                else:
                    print(f"{prefix}  ⚠️  Error: Exception in get_elements after {max_retries} attempts: {type(e).__name__}: {str(e)}")
                    import traceback
                    traceback.print_exc()
                    return 0
        
        # Should not reach here, but just in case
        return 0
    
    async def get_projects(self, session: aiohttp.ClientSession, iteration: int = 0):
        """Get all projects"""
        url = f"{self.endpoint}/projects"
        headers = {
            "Authorization": f"Bearer {self.bearer_token}"
        }
        
        start_time = time.time()
        try:
            async with session.get(
                url,
                headers=headers,
                timeout=self._client_timeout(),
            ) as response:
                duration_ms = (time.time() - start_time) * 1000
                
                try:
                    data = await response.json()
                except Exception as e:
                    error_msg = f"Failed to parse get_projects response: {str(e)}"
                    self.tracker.record_error('get_projects', iteration, 'JsonParseError', error_msg, {'status': response.status})
                    raise Exception(error_msg)
                
                project_count = len(data) if isinstance(data, list) else len(data.get('projects', []))
                self.tracker.record_metric('get_projects', duration_ms, {
                    'project_count': project_count,
                    'status': response.status
                })
                
                if response.status != 200:
                    error_msg = f"Failed to get projects: {response.status} - {data}"
                    self.tracker.record_error('get_projects', iteration, 'GetProjectsFailed', error_msg, {'status': response.status})
                    raise Exception(error_msg)
        except Exception as e:
            # Re-raise to let caller handle
            if 'get_projects' not in str(e).lower():
                error_msg = f"Failed to get projects: {str(e)}"
                self.tracker.record_error('get_projects', iteration, 'GetProjectsError', error_msg, {})
            raise
    
    async def delete_project(self, session: aiohttp.ClientSession, project_id: str):
        """Delete a project"""
        url = f"{self.endpoint}/projects/{project_id}"
        headers = {
            "Authorization": f"Bearer {self.bearer_token}"
        }
        
        start_time = time.time()
        async with session.delete(url, headers=headers) as response:
            duration_ms = (time.time() - start_time) * 1000
            
            if response.status in (200, 202, 204):
                return True
            else:
                try:
                    error_text = await response.text()
                except:
                    error_text = f"Status {response.status}"
                print(f"    ⚠️  Warning: Failed to delete project {project_id}: {response.status} - {error_text[:100]}")
                return False
    
    async def cleanup_projects(self, session: aiohttp.ClientSession):
        """Delete all projects created during the test"""
        if not self.project_ids:
            return
        
        print(f"\n{'='*60}")
        print(f"Cleaning up {len(self.project_ids)} test projects...")
        print(f"{'='*60}\n")
        
        deleted_count = 0
        failed_count = 0
        
        for project_id in self.project_ids:
            if await self.delete_project(session, project_id):
                deleted_count += 1
                print(f"  ✓ Deleted project: {project_id}")
            else:
                failed_count += 1
        
        print(f"\nCleanup complete: {deleted_count} deleted, {failed_count} failed")
        
        # Clear the project IDs list
        self.project_ids = []
    
    
    async def run_sequential_test(self, num_iterations: int):
        """Run sequential end-to-end test"""
        print(f"\n{'='*60}")
        print(f"Running SEQUENTIAL test with {num_iterations} iterations")
        print(f"{'='*60}\n")
        
        model_json = self.load_test_model()
        self.tracker.start_time = time.time()
        self.element_counts = []  # Reset element counts for this test
        
        async with aiohttp.ClientSession() as session:
            for i in range(num_iterations):
                iteration_num = i + 1
                print(f"Iteration {iteration_num}/{num_iterations}...")
                
                try:
                    # Create project
                    project_name = f"TestProject_{int(time.time())}_{i}"
                    project_id = await self.create_project(session, project_name)
                    print(f"  ✓ Created project: {project_id}")
                    
                    # Commit model
                    commit_id = await self.commit_model(session, project_id, model_json)
                    self.commit_ids.append(commit_id)
                    print(f"  ✓ Committed model (commit ID: {commit_id})")
                    
                    # Get elements using commit ID from the same project
                    # get_elements will retry automatically if commit graph is not ready yet
                    element_count = await self.get_elements(session, project_id, commit_id)
                    self.element_counts.append(element_count)  # Track element count
                    if element_count > 0:
                        print(f"  ✓ Retrieved {element_count} elements")
                        
                        # Get projects after each successful get_elements
                        await self.get_projects(session, iteration=iteration_num)
                        print(f"  ✓ Retrieved all projects")
                    else:
                        print(f"  ⚠️  Elements endpoint not available or returned 0 elements")
                except Exception as e:
                    error_type = type(e).__name__
                    error_msg = str(e)
                    print(f"  ❌ Iteration {iteration_num} failed: {error_type}: {error_msg[:200]}")
                    self.tracker.record_error(
                        operation='iteration',
                        iteration=iteration_num,
                        error_type=error_type,
                        error_message=error_msg,
                        metadata={'project_id': project_id if 'project_id' in locals() else None}
                    )
                    # Continue to next iteration instead of stopping
                    continue
            
            # Cleanup: Delete all created projects
            await self.cleanup_projects(session)
    
    async def run_concurrent_test(self, num_iterations: int, concurrency: int):
        """Run concurrent test with specified concurrency"""
        print(f"\n{'='*60}")
        print(f"Running CONCURRENT test with {num_iterations} iterations, concurrency={concurrency}")
        print(f"{'='*60}\n")
        
        model_json = self.load_test_model()
        model_size_mb = len(model_json) / (1024 * 1024)
        
        # Warn about potential gateway timeouts with large models and high concurrency
        if model_size_mb > 10 and concurrency > 2:
            print(f"⚠️  WARNING: Large model detected ({model_size_mb:.1f} MB) with high concurrency ({concurrency})")
            print(f"   This may cause 524 Gateway Timeout errors (gateway timeout is ~100 seconds)")
            print(f"   Recommendation: Reduce concurrency to 1-2 for large models, or use sequential mode")
            print(f"   Current: {concurrency} concurrent commits × {model_size_mb:.1f} MB = ~{concurrency * model_size_mb:.1f} MB total\n")
        
        self.tracker.start_time = time.time()
        
        semaphore = asyncio.Semaphore(concurrency)
        
        # Use a lock to ensure project_ids list is thread-safe
        project_ids_lock = asyncio.Lock()
        
        async def run_iteration(i: int):
            iteration_prefix = f"[Iter {i+1}]"
            project_id = None
            try:
                async with semaphore:
                    # Use context manager to ensure session is properly closed
                    async with aiohttp.ClientSession() as session:
                        print(f"{iteration_prefix} Starting iteration {i+1}...")
                        
                        # Create project
                        project_name = f"ConcurrentTest_{int(time.time())}_{i}"
                        project_id = await self.create_project(session, project_name)
                        print(f"{iteration_prefix} ✓ Created project: {project_id}")
                        
                        # Validate project_id is not None/empty
                        if not project_id:
                            raise Exception(f"Failed to get project ID from create_project response")
                        
                        # Validate project_id format (should be a UUID)
                        if len(project_id) < 30 or not isinstance(project_id, str):
                            raise Exception(f"Invalid project_id format: {project_id} (type: {type(project_id)})")
                        
                        # IMPORTANT: Add project_id to list immediately after creation (with lock)
                        # This ensures we track it even if the iteration fails later
                        async with project_ids_lock:
                            if project_id not in self.project_ids:
                                self.project_ids.append(project_id)
                        
                        # Verify project exists before committing (safety check)
                        # This helps catch cases where project was deleted by another process
                        verify_url = f"{self.endpoint}/projects/{project_id}"
                        async with session.get(
                            verify_url,
                            headers={"Authorization": f"Bearer {self.bearer_token}"},
                            timeout=self._client_timeout(),
                        ) as verify_response:
                            if verify_response.status == 404:
                                raise Exception(f"Project {project_id} does not exist (404) - cannot commit. Project may have been deleted.")
                            elif verify_response.status != 200:
                                print(f"{iteration_prefix} ⚠️  Warning: Project verification returned {verify_response.status}, but proceeding with commit")
                        
                        # Commit model using the SAME project_id from this iteration
                        # This ensures we're committing to the project we just created
                        print(f"{iteration_prefix} → Committing to project: {project_id}")
                        commit_id = await self.commit_model(session, project_id, model_json)
                        print(f"{iteration_prefix} ✓ Committed model (commit ID: {commit_id}) to project: {project_id}")
                        
                        # Get elements using commit ID from the same project
                        # get_elements will retry automatically if commit graph is not ready yet
                        element_count = await self.get_elements(session, project_id, commit_id, iteration_prefix=iteration_prefix)
                        self.element_counts.append(element_count)  # Track element count
                        
                        # Get projects after each successful get_elements
                        if element_count > 0:
                            await self.get_projects(session, iteration=i+1)
                            print(f"{iteration_prefix} ✓ Retrieved {element_count} elements and all projects")
                        else:
                            print(f"{iteration_prefix} ⚠️  Elements endpoint not available or returned 0 elements")
                        
                        # Session will be closed automatically by context manager
                        # All HTTP requests must complete before we exit this block
                        print(f"{iteration_prefix} ✓ Completed iteration {i+1}")
                        return {'success': True, 'iteration': i+1, 'project_id': project_id}
            except Exception as e:
                error_type = type(e).__name__
                error_msg = str(e)
                print(f"{iteration_prefix} ❌ Failed: {error_type}: {error_msg[:200]}")
                
                # Record error in tracker
                self.tracker.record_error(
                    operation='iteration',
                    iteration=i+1,
                    error_type=error_type,
                    error_message=error_msg,
                    metadata={'project_id': project_id}
                )
                
                # Still track the project ID for cleanup (with lock)
                # BUT: If it was a timeout, the server might still be processing
                # We'll still track it but wait longer before cleanup
                if project_id:
                    async with project_ids_lock:
                        if project_id not in self.project_ids:
                            self.project_ids.append(project_id)
                
                return {'success': False, 'iteration': i+1, 'project_id': project_id, 'error': error_msg}
        
        # Run all iterations concurrently (continue even if some fail)
        # IMPORTANT: asyncio.gather waits for ALL coroutines to complete
        # Each iteration uses 'async with session' which ensures all HTTP requests
        # complete and the session is closed before the coroutine returns
        tasks = [run_iteration(i) for i in range(num_iterations)]
        results = await asyncio.gather(*tasks, return_exceptions=True)
        
        # Small delay to ensure all HTTP connections are fully closed
        # This is a safety measure in case there are any lingering connections
        await asyncio.sleep(0.5)
        
        # Report summary
        successful = 0
        failed = 0
        for r in results:
            if isinstance(r, dict):
                if r.get('success'):
                    successful += 1
                else:
                    failed += 1
            else:
                # Unexpected exception that wasn't caught
                failed += 1
                print(f"⚠️  Unexpected error in iteration: {type(r).__name__}: {str(r)[:200]}")
        
        if failed > 0:
            print(f"\n⚠️  {failed} iteration(s) failed out of {num_iterations} total ({successful} succeeded)")
        else:
            print(f"\n✓ All {num_iterations} iterations completed successfully")
        
        # Final operations - use a NEW session for cleanup
        # All previous sessions are already closed (via context managers)
        # Cleanup happens AFTER all iterations have fully completed and their sessions are closed
        
        # IMPORTANT: Wait longer to ensure all server-side operations complete
        # Even though client requests are done, server might still be processing commits
        # When a TimeoutError occurs, the client gives up but the server continues processing
        # We need to wait long enough for those in-flight commits to complete
        print(f"\nWaiting for all server operations to complete before cleanup...")
        print(f"  (This ensures any in-flight commits from timeouts can complete)")
        await asyncio.sleep(10)  # Wait 10 seconds for any in-flight commits to complete
        
        async with aiohttp.ClientSession() as session:
            # Final get projects (no iteration number for final call)
            await self.get_projects(session, iteration=0)
            print(f"  ✓ Final: Retrieved all projects")
            
            # Cleanup: Delete all created projects
            # IMPORTANT: This only deletes projects that were successfully created
            # and tracked in self.project_ids. All iterations have completed by now.
            # We wait 5 seconds above to ensure no commits are still in flight.
            if self.project_ids:
                print(f"\nStarting cleanup of {len(self.project_ids)} project(s)...")
                await self.cleanup_projects(session)
            else:
                print(f"\nNo projects to clean up.")


def generate_charts(tracker: PerformanceTracker, output_dir: Path, test_mode: str, test_number: int = 0):
    """Generate performance charts - time per iteration and error charts"""
    # Create test-specific subdirectory: test_mode/test0, test_mode/test1, etc.
    mode_dir = output_dir / test_mode / f"test{test_number}"
    mode_dir.mkdir(parents=True, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    
    # Create 4 separate plots for each operation: create_project, commit, get_projects, get_elements
    operations = ['create_project', 'commit', 'get_projects', 'get_elements']
    operation_titles = ['Create Project', 'Commit', 'Get Projects', 'Get Elements']
    
    charts = {}
    
    # Generate individual charts for each operation
    for op, title in zip(operations, operation_titles):
        if op not in tracker.metrics or not tracker.metrics[op]:
            continue
        
        measurements = tracker.metrics[op]
        iterations = list(range(1, len(measurements) + 1))
        durations = [m['duration_ms'] for m in measurements]
        
        fig, ax = plt.subplots(figsize=(12, 6))
        ax.plot(iterations, durations, marker='o', linewidth=2, markersize=6, alpha=0.7)
        ax.set_xlabel('Iteration')
        ax.set_ylabel('Time (ms)')
        ax.set_title(f'{title} - Time per Iteration - {test_mode}')
        ax.grid(True, alpha=0.3)
        
        # Add statistics text
        if durations:
            avg_time = sum(durations) / len(durations)
            min_time = min(durations)
            max_time = max(durations)
            stats_text = f'Avg: {avg_time:.2f}ms\nMin: {min_time:.2f}ms\nMax: {max_time:.2f}ms'
            ax.text(0.02, 0.98, stats_text, transform=ax.transAxes, 
                   verticalalignment='top', bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))
        
        plt.tight_layout()
        chart_path = mode_dir / f"{op}_time_per_iteration_{test_mode}_{timestamp}.png"
        plt.savefig(chart_path, dpi=150)
        print(f"  ✓ Saved chart: {chart_path}")
        plt.close()
        
        charts[op] = chart_path
    
    # Generate error chart - multiline showing errors per iteration for each operation
    # Track errors by operation and iteration: 0 = no error, 1 = error
    operations = ['create_project', 'commit', 'get_projects', 'get_elements']
    operation_titles = ['Create Project', 'Commit', 'Get Projects', 'Get Elements']
    operation_colors = ['#FF6B6B', '#4ECDC4', '#45B7D1', '#FFA07A']  # Different colors for each operation
    
    # Find the maximum iteration number
    max_iteration = 0
    if tracker.errors:
        max_iteration = max(e.get('iteration', 0) for e in tracker.errors)
    
    # Also check metrics to get max iteration
    for op in operations:
        if op in tracker.metrics and tracker.metrics[op]:
            max_iteration = max(max_iteration, len(tracker.metrics[op]))
    
    if max_iteration > 0:
        # Initialize error matrix: operation -> iteration -> error (0 or 1)
        error_matrix = {op: [0] * max_iteration for op in operations}
        
        # Mark errors in the matrix
        if tracker.errors:
            for error in tracker.errors:
                operation = error.get('operation', '')
                iteration = error.get('iteration', 0)
                
                # Skip errors with iteration 0 (not tied to a specific iteration)
                if iteration == 0:
                    continue
                
                # Map operation names to our standard operation keys
                op_key = None
                operation_lower = operation.lower()
                if operation == 'create_project' or 'create_project' in operation_lower:
                    op_key = 'create_project'
                elif operation == 'commit' or operation_lower == 'commit':
                    op_key = 'commit'
                elif operation == 'get_projects' or 'get_projects' in operation_lower:
                    op_key = 'get_projects'
                elif operation == 'get_elements' or 'get_elements' in operation_lower:
                    op_key = 'get_elements'
                elif operation == 'iteration':
                    # For iteration-level errors, try to infer which operation failed from error message
                    error_msg = error.get('error_message', '').lower()
                    metadata = error.get('metadata', {})
                    # Try to infer from error message
                    if 'create' in error_msg or ('project' in error_msg and 'create' in error_msg):
                        op_key = 'create_project'
                    elif 'commit' in error_msg:
                        op_key = 'commit'
                    elif 'elements' in error_msg:
                        op_key = 'get_elements'
                    elif 'get projects' in error_msg or 'projects' in error_msg:
                        op_key = 'get_projects'
                    # If we can't determine, mark all operations as having an error for this iteration
                    # (This is a fallback - ideally we should know which operation failed)
                
                if op_key and op_key in error_matrix and 1 <= iteration <= max_iteration:
                    error_matrix[op_key][iteration - 1] = 1  # Mark as error (1)
        
        # Create multiline chart
        fig, ax = plt.subplots(figsize=(12, 6))
        iterations = list(range(1, max_iteration + 1))
        
        # Plot each operation as a separate line
        for op, title, color in zip(operations, operation_titles, operation_colors):
            error_values = error_matrix[op]
            ax.plot(iterations, error_values, marker='o', linewidth=2, markersize=8, 
                   label=title, color=color, alpha=0.8)
        
        ax.set_xlabel('Iteration')
        ax.set_ylabel('Error Status')
        ax.set_title(f'Errors per Iteration by Operation - {test_mode}')
        ax.set_ylim(-0.1, 1.1)
        ax.set_yticks([0, 1])
        ax.set_yticklabels(['No Error', 'Error'])
        ax.legend(loc='upper right')
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        error_chart_path = mode_dir / f"errors_per_iteration_{test_mode}_{timestamp}.png"
        plt.savefig(error_chart_path, dpi=150)
        print(f"  ✓ Saved error chart: {error_chart_path}")
        plt.close()
        charts['errors_per_iteration'] = error_chart_path
    
    return charts


def print_summary(tracker: PerformanceTracker):
    """Print performance summary"""
    print(f"\n{'='*60}")
    print("PERFORMANCE SUMMARY")
    print(f"{'='*60}\n")
    
    stats = tracker.get_statistics()
    for op, stat in stats.items():
        print(f"{op.replace('_', ' ').title()}:")
        print(f"  Count:    {stat['count']}")
        print(f"  Min:       {stat['min']:.2f} ms")
        print(f"  Max:       {stat['max']:.2f} ms")
        print(f"  Average:   {stat['avg']:.2f} ms")
        print(f"  Total:     {stat['total']:.2f} ms")
        print()
    
    # Print error summary
    error_stats = tracker.get_error_statistics()
    if error_stats['total_errors'] > 0:
        print(f"{'='*60}")
        print("ERROR SUMMARY")
        print(f"{'='*60}\n")
        print(f"Total Errors: {error_stats['total_errors']}")
        print(f"\nErrors by Operation:")
        for op, count in error_stats['errors_by_operation'].items():
            print(f"  {op}: {count}")
        print(f"\nErrors by Type:")
        for error_type, count in error_stats['errors_by_type'].items():
            print(f"  {error_type}: {count}")
        print()
    


async def main():
    # Load environment variables from .env file
    load_dotenv()
    
    parser = argparse.ArgumentParser(description='Run Flexo API performance tests')
    parser.add_argument('--mode', choices=['sequential', 'concurrent', 'both'], 
                       default='both', help='Test execution mode (default: both)')
    parser.add_argument('--iterations', type=int, default=10, 
                       help='Number of iterations (default: 10)')
    parser.add_argument('--concurrency', type=int, default=5, 
                       help='Concurrency level for concurrent mode (default: 5)')
    parser.add_argument('--endpoint', default=ENDPOINT, 
                       help=f'Flexo API endpoint (default: {ENDPOINT})')
    parser.add_argument('--token', default=None, 
                       help='Bearer token for authentication (default: from .env FLEXO_TOKEN)')
    parser.add_argument('--model', default=TEST_MODEL_FILE, 
                       help=f'Path to test model JSON file (default: {TEST_MODEL_FILE})')
    parser.add_argument('--output', default='performance_results', 
                       help='Output directory for charts and results (default: performance_results)')
    parser.add_argument('--request-timeout', type=int, default=DEFAULT_REQUEST_TIMEOUT_SECONDS,
                       help=f'Request timeout in seconds for all HTTP calls (default: {DEFAULT_REQUEST_TIMEOUT_SECONDS}s)')
    
    args = parser.parse_args()
    
    # Get token from argument or environment variable
    bearer_token = args.token or os.getenv('FLEXO_TOKEN')
    if not bearer_token:
        print("Error: Bearer token is required. Provide it via:")
        print("  1. --token argument, or")
        print("  2. FLEXO_TOKEN environment variable in .env file")
        print("\nCreate a .env file with:")
        print("  FLEXO_TOKEN=your_token_here")
        return
    
    # Verify test model exists
    model_path = Path(args.model)
    if not model_path.exists():
        print(f"Error: Test model file not found: {model_path}")
        return
    
    runner = FlexoTestRunner(
        args.endpoint,
        bearer_token,
        str(model_path),
        request_timeout=args.request_timeout,
    )
    
    print(f"\nHTTP request timeout set to {runner.request_timeout} seconds\n")
    output_dir = Path(args.output)
    output_dir.mkdir(exist_ok=True)
    
    results = {}
    
    # Find next available test number for each mode
    def get_next_test_number(mode: str) -> int:
        mode_dir = output_dir / mode
        if not mode_dir.exists():
            return 0
        existing_tests = [d.name for d in mode_dir.iterdir() if d.is_dir() and d.name.startswith('test')]
        if not existing_tests:
            return 0
        test_numbers = []
        for test_dir in existing_tests:
            try:
                num = int(test_dir.replace('test', ''))
                test_numbers.append(num)
            except ValueError:
                continue
        return max(test_numbers) + 1 if test_numbers else 0
    
    try:
        if args.mode in ('sequential', 'both'):
            test_num = get_next_test_number('sequential')
            print(f"\nUsing test directory: sequential/test{test_num}")
            
            await runner.run_sequential_test(args.iterations)
            print_summary(runner.tracker)
            
            # Create test-specific subdirectory
            test_dir = output_dir / 'sequential' / f'test{test_num}'
            test_dir.mkdir(parents=True, exist_ok=True)
            
            charts = generate_charts(runner.tracker, output_dir, 'sequential', test_num)
            results['sequential'] = {
                'stats': runner.tracker.get_statistics(),
                'error_stats': runner.tracker.get_error_statistics(),
                'charts': charts
            }
            
            # Save results to JSON in test-specific subdirectory
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            results_file = test_dir / f"sequential_results_{timestamp}.json"
            
            # Prepare detailed results with per-iteration data
            detailed_results = {
                'mode': 'sequential',
                'test_number': test_num,
                'iterations': args.iterations,
                'statistics': runner.tracker.get_statistics(),
                'error_statistics': runner.tracker.get_error_statistics(),
                'projects_created': len(runner.project_ids),
                'projects_cleaned_up': True,
                'timestamp': datetime.now().isoformat()
            }
            
            with open(results_file, 'w') as f:
                json.dump(detailed_results, f, indent=2)
            print(f"  ✓ Saved results: {results_file}")
            
            # Save error log if there are errors
            if runner.tracker.errors:
                error_log_file = test_dir / f"error_log_{timestamp}.json"
                with open(error_log_file, 'w') as f:
                    json.dump(runner.tracker.get_error_statistics(), f, indent=2)
                print(f"  ✓ Saved error log: {error_log_file}")
        
        if args.mode in ('concurrent', 'both'):
            # Create new tracker for concurrent test
            runner.tracker = PerformanceTracker()
            runner.project_ids = []  # CRITICAL: Clear project IDs from previous runs
            runner.commit_ids = []  # Clear commit IDs from previous runs
            runner.element_counts = []  # Reset element counts for concurrent test
            
            test_num = get_next_test_number('concurrent')
            print(f"\nUsing test directory: concurrent/test{test_num}")
            
            await runner.run_concurrent_test(args.iterations, args.concurrency)
            print_summary(runner.tracker)
            
            # Create test-specific subdirectory
            test_dir = output_dir / 'concurrent' / f'test{test_num}'
            test_dir.mkdir(parents=True, exist_ok=True)
            
            charts = generate_charts(runner.tracker, output_dir, 'concurrent', test_num)
            results['concurrent'] = {
                'stats': runner.tracker.get_statistics(),
                'error_stats': runner.tracker.get_error_statistics(),
                'charts': charts
            }
            
            # Save results to JSON in test-specific subdirectory
            timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
            results_file = test_dir / f"concurrent_results_{timestamp}.json"
            
            # Prepare detailed results with per-iteration data
            detailed_results = {
                'mode': 'concurrent',
                'test_number': test_num,
                'iterations': args.iterations,
                'concurrency': args.concurrency,
                'statistics': runner.tracker.get_statistics(),
                'error_statistics': runner.tracker.get_error_statistics(),
                'projects_created': len(runner.project_ids),
                'projects_cleaned_up': True,
                'timestamp': datetime.now().isoformat()
            }
            
            with open(results_file, 'w') as f:
                json.dump(detailed_results, f, indent=2)
            print(f"  ✓ Saved results: {results_file}")
            
            # Save error log if there are errors
            if runner.tracker.errors:
                error_log_file = test_dir / f"error_log_{timestamp}.json"
                with open(error_log_file, 'w') as f:
                    json.dump(runner.tracker.get_error_statistics(), f, indent=2)
                print(f"  ✓ Saved error log: {error_log_file}")
    except Exception as e:
        print(f"\n⚠️  Error during test execution: {e}")
        print("Attempting to cleanup projects before exiting...")
        # Try to cleanup even if test failed
        try:
            async with aiohttp.ClientSession() as session:
                await runner.cleanup_projects(session)
        except Exception as cleanup_error:
            print(f"⚠️  Cleanup also failed: {cleanup_error}")
        raise
    
    print(f"\n{'='*60}")
    print("All tests completed!")
    print(f"Results saved to: {output_dir}")
    print(f"{'='*60}\n")


if __name__ == '__main__':
    asyncio.run(main())