test_data_loading.py

#!/usr/bin/env python3
"""
Test script to verify data loading functionality
"""

import sys
import pandas as pd
import yfinance as yf
import datetime
import logging

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

def test_stock_data_loading():
    """Test the stock data loading functionality"""
    print("🧪 Testing Stock Data Loading...")
    
    # Test parameters
    test_ticker = "AAPL"
    years_back = 2
    frequency = '1d'
    
    try:
        # Use dynamic date range
        end_date = datetime.datetime.now()
        start_date = datetime.datetime(end_date.year - years_back, end_date.month, end_date.day)
        
        # Format dates for yfinance
        start = start_date.strftime('%Y-%m-%d')
        end = end_date.strftime('%Y-%m-%d')
        
        print(f"📅 Fetching {test_ticker} data from {start} to {end} with {frequency} frequency...")
        
        # Download data
        df = yf.download(test_ticker, start=start, end=end, interval=frequency)
        
        if df.empty:
            print(f"❌ No data returned for {test_ticker}")
            return False
            
        print(f"✅ Successfully fetched {test_ticker} data:")
        print(f"   📊 Shape: {df.shape}")
        print(f"   📈 Date range: {df.index.min()} to {df.index.max()}")
        print(f"   💰 Price range: ${df['Close'].min():.2f} - ${df['Close'].max():.2f}")
        
        # Process data like in the app
        df = df.reset_index()[['Date', 'Close']]
        df = df.rename(columns={'Close': test_ticker})
        
        print(f"   🔄 Processed data shape: {df.shape}")
        print(f"   📋 Columns: {list(df.columns)}")
        print(f"   🎯 Data points: {len(df):,}")
        
        # Show sample data
        print("\n📋 Sample data:")
        print(df.head(3).to_string(index=False))
        print("...")
        print(df.tail(3).to_string(index=False))
        
        return True
        
    except Exception as e:
        print(f"❌ Error fetching data for {test_ticker}: {e}")
        return False

def test_multiple_stocks():
    """Test loading multiple stocks"""
    print("\n🧪 Testing Multiple Stock Loading...")
    
    test_tickers = ["AAPL", "MSFT", "GOOGL"]
    years_back = 1
    frequency = '1d'
    
    results = {}
    
    for ticker in test_tickers:
        try:
            end_date = datetime.datetime.now()
            start_date = datetime.datetime(end_date.year - years_back, end_date.month, end_date.day)
            
            start = start_date.strftime('%Y-%m-%d')
            end = end_date.strftime('%Y-%m-%d')
            
            df = yf.download(ticker, start=start, end=end, interval=frequency)
            
            if not df.empty:
                df = df.reset_index()[['Date', 'Close']]
                df = df.rename(columns={'Close': ticker})
                results[ticker] = df
                print(f"✅ {ticker}: {len(df):,} data points")
            else:
                print(f"❌ {ticker}: No data")
                
        except Exception as e:
            print(f"❌ {ticker}: Error - {e}")
    
    if results:
        total_points = sum(len(df) for df in results.values())
        print(f"\n📊 Summary:")
        print(f"   ✅ Successfully loaded: {len(results)}/{len(test_tickers)} stocks")
        print(f"   🎯 Total data points: {total_points:,}")
        
        # Test different frequencies
        print(f"\n🔄 Testing different frequencies for AAPL...")
        frequencies = ['1d', '1wk', '1mo']
        
        for freq in frequencies:
            try:
                df = yf.download("AAPL", start=start, end=end, interval=freq)
                if not df.empty:
                    print(f"   {freq}: {len(df):,} data points")
                else:
                    print(f"   {freq}: No data")
            except Exception as e:
                print(f"   {freq}: Error - {e}")
        
        return True
    else:
        print("❌ Failed to load any stock data")
        return False

def test_sp500_file():
    """Test loading S&P 500 characteristics file"""
    print("\n🧪 Testing S&P 500 File Loading...")
    
    try:
        sp = pd.read_csv('SP_500/Characteristics.csv', index_col='Ticker')
        print(f"✅ Successfully loaded S&P 500 data:")
        print(f"   📊 Shape: {sp.shape}")
        print(f"   📋 Columns: {list(sp.columns)}")
        
        # Check sectors
        sectors = sp.Sector.dropna().unique()
        print(f"   🏢 Sectors: {len(sectors)}")
        print(f"   📈 Sample tickers: {sp.index[:10].tolist()}")
        
        return True
        
    except Exception as e:
        print(f"❌ Error loading S&P 500 file: {e}")
        return False

if __name__ == "__main__":
    print("🚀 Starting Data Loading Tests...\n")
    
    # Run tests
    test1 = test_sp500_file()
    test2 = test_stock_data_loading()
    test3 = test_multiple_stocks()
    
    # Summary
    print(f"\n📋 Test Results:")
    print(f"   S&P 500 File: {'✅ PASS' if test1 else '❌ FAIL'}")
    print(f"   Single Stock: {'✅ PASS' if test2 else '❌ FAIL'}")
    print(f"   Multiple Stocks: {'✅ PASS' if test3 else '❌ FAIL'}")
    
    if all([test1, test2, test3]):
        print(f"\n🎉 All tests passed! The dashboard should work correctly.")
        sys.exit(0)
    else:
        print(f"\n⚠️  Some tests failed. Check the errors above.")
        sys.exit(1)