test_run.py•4.11 kB
# test_run.py - Simple test of the enhanced system
import sys
import warnings
warnings.filterwarnings('ignore')
# Simple test with fallback to basic functionality
def simple_test():
"""Test the enhanced system with basic functionality."""
try:
# Try to import enhanced modules
from models.arima_model import get_arima_forecast
from models.hybrid_model import train_xgboost_on_residuals
print("Testing Enhanced Trading System")
print("=" * 40)
# Test with one ticker
ticker = 'AAPL'
print(f"\nTesting with {ticker}...")
# Get ARIMA forecast
print("1. Running ARIMA analysis...")
arima_forecast, residuals, df = get_arima_forecast(ticker)
print(f" ARIMA forecast: ${arima_forecast:.2f}")
# Train XGBoost
print("2. Training XGBoost on residuals...")
xgb_model, latest_features = train_xgboost_on_residuals(residuals, df)
# Get hybrid forecast
residual_correction = xgb_model.predict(latest_features)[0]
hybrid_forecast = arima_forecast + residual_correction
last_price = df['Close'].iloc[-1]
change = (hybrid_forecast - last_price) / last_price
print("3. Results:")
print(f" Last Price: ${last_price:.2f}")
print(f" ARIMA Forecast: ${arima_forecast:.2f}")
print(f" Hybrid Forecast: ${hybrid_forecast:.2f}")
print(f" Expected Change: {change:+.2%}")
# Generate signal
if change > 0.02:
signal = "BUY"
elif change < -0.02:
signal = "SELL"
else:
signal = "HOLD"
print(f" Trading Signal: {signal}")
print("\n✅ Test completed successfully!")
return True
except ImportError as e:
print(f"❌ Import error: {e}")
print("Make sure you have the enhanced model files in the models/ directory")
return False
except Exception as e:
print(f"❌ Error during analysis: {e}")
return False
def full_test():
"""Test the full enhanced system."""
try:
# Import the enhanced main module
from main import TradingSignalGenerator, simple_analysis
print("Testing Full Enhanced System")
print("=" * 40)
# Test simple analysis first
print("\n1. Testing simple analysis...")
simple_analysis()
print("\n2. Testing enhanced analysis...")
# Initialize with minimal config
generator = TradingSignalGenerator(
save_results=False, # Don't save files for test
parallel_processing=False # Sequential for testing
)
# Test with just 2 tickers
test_tickers = ['AAPL', 'MSFT']
generator.config['tickers'] = test_tickers
# Run analysis
signals = generator.analyze_portfolio(test_tickers)
# Print summary
generator.print_summary(signals)
print("\n✅ Full test completed successfully!")
return True
except ImportError as e:
print(f"❌ Import error: {e}")
print("Make sure you have the enhanced main.py file")
return False
except Exception as e:
print(f"❌ Error during full test: {e}")
return False
if __name__ == "__main__":
print("Enhanced Trading System Test")
print("=" * 50)
# Test basic functionality first
if simple_test():
print("\n" + "="*50)
# If basic test passes, try full system
full_test()
else:
print("\n❌ Basic test failed. Please check your setup.")
print("\nSetup checklist:")
print("1. Install required packages: pip install yfinance pmdarima xgboost scikit-learn")
print("2. Create models/ directory with __init__.py")
print("3. Copy enhanced model code to models/arima_model.py and models/hybrid_model.py")
print("4. Copy enhanced main code to main.py")