🔍 Code Extractor

A test function that validates the cleanup functionality of virtual environments in project directories by testing on a specific session, measuring disk space before/after cleanup, and verifying that important files are preserved.

Purpose

This function serves as a comprehensive integration test for the venv cleanup feature. It tests the AgentExecutor's cleanup_venv_directories method on a hardcoded test session ID, measuring disk space usage before and after cleanup, verifying that virtual environments are removed while preserving critical project files (analysis_results.json, analysis_script.py). The function provides detailed console output showing the cleanup process and results.

Source Code

def main():
    # Initialize config and agent executor
    config = Config()
    agent_executor = AgentExecutor(config)
    
    # Test cleanup on oldest session
    test_session_id = "558edb65-de39-403f-85c3-06ebfe8fa252"
    
    print(f"\n{'='*60}")
    print(f"Testing cleanup on session: {test_session_id}")
    print(f"{'='*60}\n")
    
    # Check size before cleanup
    session_dir = Path(config.GENERATED_SCRIPTS_FOLDER) / test_session_id  # Use GENERATED_SCRIPTS_FOLDER
    if not session_dir.exists():
        print(f"ERROR: Session directory not found: {session_dir}")
        return 1
    
    # Get size before
    import subprocess
    result = subprocess.run(['du', '-sh', str(session_dir)], capture_output=True, text=True)
    size_before = result.stdout.split()[0]
    print(f"Session size BEFORE cleanup: {size_before}")
    
    # List projects
    projects = [d for d in session_dir.iterdir() if d.is_dir() and d.name.startswith('project_')]
    print(f"Found {len(projects)} project(s) in session:\n")
    
    for proj in projects:
        proj_result = subprocess.run(['du', '-sh', str(proj)], capture_output=True, text=True)
        proj_size = proj_result.stdout.split()[0]
        venv_path = proj / 'venv'
        has_venv = '✓' if venv_path.exists() else '✗'
        print(f"  {proj.name}: {proj_size} [venv: {has_venv}]")
    
    # Run cleanup
    print(f"\nRunning cleanup...")
    cleanup_result = agent_executor.cleanup_venv_directories(test_session_id)
    
    if cleanup_result['success']:
        print(f"\n✓ Cleanup completed successfully!")
        print(f"  - Cleaned venvs: {cleanup_result['cleaned_count']}")
        print(f"  - Space freed: {cleanup_result['space_freed_mb']:.2f} MB")
        print(f"  - Preserved projects: {cleanup_result['preserved_count']}")
    else:
        print(f"\n✗ Cleanup failed: {cleanup_result.get('error', 'Unknown error')}")
        return 1
    
    # Get size after
    result = subprocess.run(['du', '-sh', str(session_dir)], capture_output=True, text=True)
    size_after = result.stdout.split()[0]
    print(f"\nSession size AFTER cleanup: {size_after}")
    
    # Verify venvs are gone but projects remain
    print(f"\nVerifying cleanup:\n")
    for proj in projects:
        proj_result = subprocess.run(['du', '-sh', str(proj)], capture_output=True, text=True)
        proj_size = proj_result.stdout.split()[0]
        venv_path = proj / 'venv'
        has_venv = '✓' if venv_path.exists() else '✗'
        print(f"  {proj.name}: {proj_size} [venv: {has_venv}]")
        
        # Check that important files are preserved
        analysis_file = proj / 'analysis_results.json'
        script_file = proj / 'analysis_script.py'
        if analysis_file.exists():
            print(f"    ✓ analysis_results.json preserved")
        if script_file.exists():
            print(f"    ✓ analysis_script.py preserved")
    
    print(f"\n{'='*60}")
    print(f"Cleanup test completed successfully!")
    print(f"{'='*60}\n")
    
    return 0

Return Value

Returns an integer exit code: 0 for successful cleanup test completion, 1 if the session directory is not found or if cleanup fails. This follows standard Unix exit code conventions where 0 indicates success.

Dependencies

• pathlib
• subprocess
• smartstat_config
• agent_executor

Required Imports

from pathlib import Path
from smartstat_config import Config
from agent_executor import AgentExecutor
import subprocess

Usage Example

if __name__ == '__main__':
    import sys
    from pathlib import Path
    from smartstat_config import Config
    from agent_executor import AgentExecutor
    import subprocess
    
    exit_code = main()
    sys.exit(exit_code)

Best Practices

• This function is hardcoded to test a specific session ID and should be modified or parameterized for production use
• Requires Unix-like system with 'du' command; will not work on Windows without modification
• Should be run with appropriate file system permissions to access and clean session directories
• The function performs destructive operations (deleting venv directories) and should only be run on test data
• Consider adding error handling for subprocess calls that might fail on different systems
• The test session ID should exist before running this function, or the function will return early with error code 1
• Output is verbose and designed for manual inspection; consider adding structured logging for automated testing

Tags

Similar Components

AI-powered semantic similarity - components with related functionality:

• function main_v66 80.2% similar

  A cleanup utility function that removes virtual environment directories from all SmartStat sessions while preserving project files, reporting disk space freed and cleanup statistics.

  From: /tf/active/vicechatdev/vice_ai/run_full_cleanup.py

• function main_v25 59.3% similar

  Orchestrates and executes a comprehensive test suite for the Vice AI Data Analysis Integration, running multiple test functions, creating test datasets, and providing detailed pass/fail reporting.

  From: /tf/active/vicechatdev/vice_ai/test_integration.py

• function main_v22 57.3% similar

  Orchestrates and executes a comprehensive test suite for a Contract Validity Analyzer system, running tests for configuration, FileCloud connection, document processing, LLM client, and full analyzer functionality.

  From: /tf/active/vicechatdev/contract_validity_analyzer/test_implementation.py

• function check_virtual_env 56.6% similar

  Detects whether the Python interpreter is running inside a virtual environment by checking system attributes.

  From: /tf/active/vicechatdev/email-forwarder/run_service.py

• function main_v40 56.4% similar

  Test orchestration function that executes a comprehensive test suite for DocChat's multi-LLM model selection feature and reports results.

  From: /tf/active/vicechatdev/docchat/test_model_selection.py