test_api.py

"""
IBKR API Test Suite - FIXED VERSION
-------------------
Comprehensive tests for all API endpoints.

Before running:
1. Make sure IB Gateway or TWS is running on port 4002
2. Install test dependencies: pip install requests
3. Start the API server: uvicorn main:app --reload
4. Run tests: python test_api_fixed.py -v
"""

import requests
import time
from typing import Dict, Any

# Configuration
BASE_URL = "http://localhost:8000"
TEST_SYMBOL = "AAPL"  # Change if needed
TEST_QUANTITY = 1


class Colors:
    """ANSI color codes for prettier output"""
    GREEN = '\033[92m'
    RED = '\033[91m'
    YELLOW = '\033[93m'
    BLUE = '\033[94m'
    RESET = '\033[0m'
    BOLD = '\033[1m'


def print_test_header(test_name: str):
    """Print a formatted test header"""
    print(f"\n{Colors.BLUE}{Colors.BOLD}{'=' * 60}{Colors.RESET}")
    print(f"{Colors.BLUE}{Colors.BOLD}TEST: {test_name}{Colors.RESET}")
    print(f"{Colors.BLUE}{Colors.BOLD}{'=' * 60}{Colors.RESET}")


def print_success(message: str):
    """Print success message"""
    print(f"{Colors.GREEN}✅ {message}{Colors.RESET}")


def print_error(message: str):
    """Print error message"""
    print(f"{Colors.RED}❌ {message}{Colors.RESET}")


def print_warning(message: str):
    """Print warning message"""
    print(f"{Colors.YELLOW}⚠️  {message}{Colors.RESET}")


def print_info(message: str):
    """Print info message"""
    print(f"{Colors.BLUE}ℹ️  {message}{Colors.RESET}")


def print_response(data: Any):
    """Pretty print response data"""
    import json
    print(f"{Colors.YELLOW}Response:{Colors.RESET}")
    print(json.dumps(data, indent=2, ensure_ascii=False))


# ============================================================================
# TEST 1: Server Health Check
# ============================================================================
def test_server_health():
    """Test if the API server is running"""
    print_test_header("Server Health Check")
    
    try:
        response = requests.get(f"{BASE_URL}/", timeout=5)
        
        if response.status_code == 200:
            print_success("Server is running")
            print_response(response.json())
            return True, response.json()
        else:
            print_error(f"Server returned status code: {response.status_code}")
            return False, None
            
    except requests.exceptions.ConnectionError:
        print_error("Cannot connect to server. Is it running?")
        print_info("Start server with: uvicorn main:app --reload")
        return False, None
    except Exception as e:
        print_error(f"Unexpected error: {e}")
        return False, None


# ============================================================================
# TEST 2: Get Account Information
# ============================================================================
def test_get_account():
    """Test GET /account endpoint"""
    print_test_header("Get Account Information")
    
    try:
        response = requests.get(f"{BASE_URL}/account", timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            print_success("Account info retrieved")
            print_response(data)
            
            # Validate response structure
            if "balance" in data and "cash" in data:
                print_success("Response structure is correct")
                return True, data
            else:
                print_error("Response missing required fields")
                return False, None
        else:
            print_error(f"Request failed with status: {response.status_code}")
            print_response(response.text)
            return False, None
            
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


# ============================================================================
# TEST 3: Get Open Positions
# ============================================================================
def test_get_positions():
    """Test GET /positions endpoint"""
    print_test_header("Get Open Positions")
    
    try:
        response = requests.get(f"{BASE_URL}/positions", timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            print_success("Positions retrieved")
            print_response(data)
            
            # Validate response structure
            if "positions" in data:
                positions = data["positions"]
                print_info(f"Found {len(positions)} open position(s)")
                
                # Validate each position structure
                for pos in positions:
                    required_fields = ["symbol", "qty", "avg_cost", "market_value"]
                    if all(field in pos for field in required_fields):
                        print_success(f"Position {pos['symbol']}: Valid structure")
                    else:
                        print_error(f"Position missing required fields")
                        return False, None
                
                return True, data
            else:
                print_error("Response missing 'positions' field")
                return False, None
        else:
            print_error(f"Request failed with status: {response.status_code}")
            return False, None
            
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


# ============================================================================
# TEST 4: Get Pending Orders
# ============================================================================
def test_get_orders():
    """Test GET /orders endpoint"""
    print_test_header("Get Pending Orders")
    
    try:
        response = requests.get(f"{BASE_URL}/orders", timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            print_success("Orders retrieved")
            print_response(data)
            
            if "orders" in data:
                orders = data["orders"]
                print_info(f"Found {len(orders)} pending order(s)")
                return True, data
            else:
                print_error("Response missing 'orders' field")
                return False, None
        else:
            print_error(f"Request failed with status: {response.status_code}")
            return False, None
            
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


# ============================================================================
# TEST 5: Get Executions
# ============================================================================
def test_get_executions():
    """Test GET /executions endpoint"""
    print_test_header("Get Executions (All)")
    
    try:
        # Test without symbol filter
        response = requests.get(f"{BASE_URL}/executions", timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            print_success("Executions retrieved")
            print_response(data)
            
            if "executions" in data:
                executions = data["executions"]
                print_info(f"Found {len(executions)} execution(s)")
                return True, data
            else:
                print_error("Response missing 'executions' field")
                return False, None
        elif response.status_code == 404:
            print_warning("Endpoint /executions not found - might not be registered")
            print_info("Check that router is properly included in main.py")
            return False, None
        else:
            print_error(f"Request failed with status: {response.status_code}")
            return False, None
            
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


def test_get_executions_filtered():
    """Test GET /executions with symbol filter"""
    print_test_header(f"Get Executions (Filtered by {TEST_SYMBOL})")
    
    try:
        response = requests.get(
            f"{BASE_URL}/executions",
            params={"symbol": TEST_SYMBOL},
            timeout=10
        )
        
        if response.status_code == 200:
            data = response.json()
            print_success(f"Executions for {TEST_SYMBOL} retrieved")
            print_response(data)
            
            if "executions" in data:
                executions = data["executions"]
                print_info(f"Found {len(executions)} execution(s) for {TEST_SYMBOL}")
                return True, data
            else:
                print_error("Response missing 'executions' field")
                return False, None
        elif response.status_code == 404:
            print_warning("Endpoint /executions not found")
            return False, None
        else:
            print_error(f"Request failed with status: {response.status_code}")
            return False, None
            
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


# ============================================================================
# TEST 6: Place Limit Order (SIMULATION MODE)
# ============================================================================
def test_place_limit_order(execute: bool = False):
    """
    Test POST /order/limit endpoint
    
    Args:
        execute: If True, actually place the order. If False, just show what would be sent.
    """
    print_test_header("Place Limit Order")
    
    if not execute:
        print_warning("SIMULATION MODE - Order will NOT be placed")
        print_info("To actually place orders, run with execute=True")
        order_data = {
            "symbol": TEST_SYMBOL,
            "action": "BUY",
            "quantity": TEST_QUANTITY,
            "limit_price": 100.0,  # Intentionally low price to avoid accidental fill
            "order_ref": "test_limit_order"
        }
        print_info("Order that would be placed:")
        print_response(order_data)
        return True, None
    
    print_warning("⚠️  LIVE ORDER MODE - This will place a REAL order!")
    print_info("Waiting 3 seconds... Press Ctrl+C to cancel")
    time.sleep(3)
    
    try:
        order_data = {
            "symbol": TEST_SYMBOL,
            "action": "BUY",
            "quantity": TEST_QUANTITY,
            "limit_price": 100.0,  # Intentionally low to avoid fill
            "order_ref": "test_limit_order"
        }
        
        response = requests.post(
            f"{BASE_URL}/order/limit",
            json=order_data,
            timeout=10
        )
        
        if response.status_code == 200:
            data = response.json()
            print_success("Limit order placed successfully")
            print_response(data)
            return True, data
        else:
            print_error(f"Request failed with status: {response.status_code}")
            print_response(response.text)
            return False, None
            
    except KeyboardInterrupt:
        print_warning("\nOrder placement cancelled by user")
        return False, None
    except Exception as e:
        print_error(f"Error: {e}")
        return False, None


# ============================================================================
# MAIN TEST RUNNER
# ============================================================================
def run_all_tests(place_real_orders: bool = False):
    """
    Run all API tests
    
    Args:
        place_real_orders: If True, will place real orders (DANGEROUS!)
    """
    print(f"\n{Colors.BOLD}{Colors.BLUE}")
    print("╔════════════════════════════════════════════════════════════╗")
    print("║          IBKR API COMPREHENSIVE TEST SUITE                 ║")
    print("╚════════════════════════════════════════════════════════════╝")
    print(f"{Colors.RESET}\n")
    
    if place_real_orders:
        print_warning("⚠️  REAL ORDER MODE ENABLED - Orders will be placed!")
    else:
        print_info("🛡️  SAFE MODE - No real orders will be placed")
    
    results = []
    
    # Test 1: Server Health
    print_info("Running test 1/6...")
    success, _ = test_server_health()
    results.append(("Server Health", success))
    if not success:
        print_error("Server is not responding. Stopping tests.")
        return
    
    time.sleep(1)
    
    # Test 2: Account Info
    print_info("Running test 2/6...")
    success, account_data = test_get_account()
    results.append(("Account Info", success))
    time.sleep(1)
    
    # Test 3: Positions
    print_info("Running test 3/6...")
    success, positions_data = test_get_positions()
    results.append(("Get Positions", success))
    time.sleep(1)
    
    # Test 4: Orders
    print_info("Running test 4/6...")
    success, orders_data = test_get_orders()
    results.append(("Get Orders", success))
    time.sleep(1)
    
    # Test 5: Executions (All)
    print_info("Running test 5/6...")
    success, exec_data = test_get_executions()
    results.append(("Get Executions", success))
    time.sleep(1)
    
    # Test 6: Limit Order
    print_info("Running test 6/6...")
    success, limit_order = test_place_limit_order(execute=place_real_orders)
    results.append(("Place Limit Order", success))
    
    # Print summary
    print(f"\n{Colors.BOLD}{Colors.BLUE}")
    print("╔════════════════════════════════════════════════════════════╗")
    print("║                      TEST SUMMARY                          ║")
    print("╚════════════════════════════════════════════════════════════╝")
    print(f"{Colors.RESET}\n")
    
    passed = sum(1 for _, success in results if success)
    total = len(results)
    
    for test_name, success in results:
        status = f"{Colors.GREEN}✅ PASS{Colors.RESET}" if success else f"{Colors.RED}❌ FAIL{Colors.RESET}"
        print(f"{status}  {test_name}")
    
    print(f"\n{Colors.BOLD}Total: {passed}/{total} tests passed{Colors.RESET}")
    
    if passed == total:
        print(f"{Colors.GREEN}{Colors.BOLD}🎉 All tests passed!{Colors.RESET}\n")
    else:
        print(f"{Colors.YELLOW}{Colors.BOLD}⚠️  Some tests failed - check details above{Colors.RESET}\n")


# ============================================================================
# SIMPLE MENU
# ============================================================================
def show_menu():
    """Display simple test menu"""
    while True:
        print(f"\n{Colors.BOLD}{Colors.BLUE}╔═══════════════════════════════════════════╗{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.BLUE}║        IBKR API TEST MENU                 ║{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.BLUE}╚═══════════════════════════════════════════╝{Colors.RESET}\n")
        
        print("1. Run All Tests (Safe Mode)")
        print("2. Test Server Health")
        print("3. Test Get Account")
        print("4. Test Get Positions")
        print("5. Test Get Orders")
        print("6. Test Get Executions")
        print("0. Exit")
        
        choice = input(f"\n{Colors.YELLOW}Enter your choice: {Colors.RESET}")
        
        if choice == "0":
            print(f"{Colors.GREEN}Goodbye!{Colors.RESET}")
            break
        elif choice == "1":
            run_all_tests(place_real_orders=False)
        elif choice == "2":
            test_server_health()
        elif choice == "3":
            test_get_account()
        elif choice == "4":
            test_get_positions()
        elif choice == "5":
            test_get_orders()
        elif choice == "6":
            test_get_executions()
            test_get_executions_filtered()
        else:
            print_error("Invalid choice")
        
        input(f"\n{Colors.BLUE}Press Enter to continue...{Colors.RESET}")


# ============================================================================
# ENTRY POINT
# ============================================================================
if __name__ == "__main__":
    import sys
    
    print(f"{Colors.BOLD}IBKR API Test Suite{Colors.RESET}\n")
    print("Make sure:")
    print("1. IB Gateway/TWS is running on port 4002")
    print("2. API server is running: uvicorn main:app --reload")
    print("3. You have test dependencies: pip install requests\n")
    
    if len(sys.argv) > 1:
        if sys.argv[1] == "--all":
            run_all_tests(place_real_orders=False)
        else:
            print(f"Usage: python test_api_fixed.py [--all]")
            print("  --all   Run all tests in safe mode")
            print("  (no args) Show interactive menu")
    else:
        show_menu()