test_code_execution.py

#!/usr/bin/env python3
"""
Test script to verify the code execution fix
"""
import requests
import json
import pandas as pd

def test_code_execution():
    """Test the code execution with a simple standardization task"""
    
    # Create sample data
    data = {
        'A': [1, 2, 3, 4, 5],
        'B': [10, 20, 30, 40, 50],
        'C': [100, 200, 300, 400, 500]
    }
    df = pd.DataFrame(data)
    
    # Convert to CSV
    csv_data = df.to_csv(index=False)
    
    # Test request
    request_data = {
        "user_prompt": "Standardize all numeric columns to have mean 0 and standard deviation 1",
        "csv_data": csv_data,
        "filename": "test_data.csv"
    }
    
    try:
        response = requests.post(
            "http://localhost:8000/process",
            json=request_data,
            timeout=30
        )
        
        if response.status_code == 200:
            result = response.json()
            print("✅ Request successful!")
            print(f"Response: {result.get('response', 'No response')[:200]}...")
            
            if result.get('updated_df'):
                print("✅ DataFrame was updated successfully")
            else:
                print("❌ DataFrame was not updated")
                
            if result.get('code'):
                print(f"Generated code:\n{result.get('code')}")
            else:
                print("❌ No code was generated")
                
        else:
            print(f"❌ Request failed with status {response.status_code}")
            print(f"Error: {response.text}")
            
    except requests.exceptions.RequestException as e:
        print(f"❌ Request failed: {e}")

if __name__ == "__main__":
    print("🧪 Testing code execution fix...")
    test_code_execution()