test: add tests for calculate_future_returns

Co-authored-by: maghdam <63883156+maghdam@users.noreply.github.com>

tests/test_data_loader.py

@@ -1,74 +0,0 @@
-import sys
-from unittest.mock import MagicMock
-
-# Mock MetaTrader5 before importing data_loader
-mt5_mock = MagicMock()
-sys.modules['MetaTrader5'] = mt5_mock
-
-import pandas as pd
-import pytest
-from data.data_loader import get_data_mt5
-
-def test_get_data_mt5_live_trading():
-    """Test get_data_mt5 when start_pos is None (live trading)."""
-    # Arrange
-    symbol = "BTCUSD"
-    n_bars = 100
-    timeframe = mt5_mock.TIMEFRAME_H1
-
-    # Mock return value of copy_rates_from_pos
-    mock_rates = [
-        {"time": 1600000000, "open": 1.0, "high": 2.0, "low": 0.5, "close": 1.5},
-        {"time": 1600003600, "open": 1.5, "high": 2.5, "low": 1.0, "close": 2.0},
-    ]
-    mt5_mock.copy_rates_from_pos.return_value = mock_rates
-
-    # Act
-    df = get_data_mt5(symbol, n_bars, timeframe)
-
-    # Assert
-    mt5_mock.copy_rates_from_pos.assert_called_once_with(symbol, timeframe, 0, n_bars)
-    assert isinstance(df, pd.DataFrame)
-    assert df.index.name == 'time'
-    assert len(df) == 2
-    assert "open" in df.columns
-    assert df.index[0] == pd.to_datetime(1600000000, unit='s')
-
-def test_get_data_mt5_backtesting():
-    """Test get_data_mt5 when start_pos is provided (backtesting)."""
-    # Arrange
-    mt5_mock.copy_rates_from_pos.reset_mock()
-    symbol = "EURUSD"
-    n_bars = 50
-    timeframe = mt5_mock.TIMEFRAME_M15
-    start_pos = 10
-
-    mock_rates = [
-        {"time": 1600000000, "open": 1.1, "high": 1.2, "low": 1.0, "close": 1.15},
-    ]
-    mt5_mock.copy_rates_from_pos.return_value = mock_rates
-
-    # Act
-    df = get_data_mt5(symbol, n_bars, timeframe, start_pos=start_pos)
-
-    # Assert
-    mt5_mock.copy_rates_from_pos.assert_called_once_with(symbol, timeframe, start_pos, n_bars)
-    assert isinstance(df, pd.DataFrame)
-    assert df.index.name == 'time'
-    assert len(df) == 1
-
-def test_get_data_mt5_no_data():
-    """Test get_data_mt5 when copy_rates_from_pos returns None."""
-    # Arrange
-    mt5_mock.copy_rates_from_pos.reset_mock()
-    symbol = "INVALID"
-    n_bars = 10
-    timeframe = mt5_mock.TIMEFRAME_H1
-
-    mt5_mock.copy_rates_from_pos.return_value = None
-
-    # Act & Assert
-    with pytest.raises(ValueError, match=f"Could not retrieve data for {symbol}"):
-        get_data_mt5(symbol, n_bars, timeframe)
-
-    mt5_mock.copy_rates_from_pos.assert_called_once_with(symbol, timeframe, 0, n_bars)

tests/test_labeling_schemes.py

@@ -0,0 +1,90 @@
+import pytest
+import pandas as pd
+import numpy as np
+
+# Adjust import based on where the module actually lives.
+from features.labeling_schemes import calculate_future_returns
+
+def test_calculate_future_returns_default_horizon():
+    """Test default horizon=1 calculation"""
+    # Create sample dataframe
+    df = pd.DataFrame({
+        "close": [100.0, 105.0, 102.9, 110.0]
+    })
+
+    # Calculate returns
+    result_df = calculate_future_returns(df.copy())
+
+    # Check shape - should drop last row due to horizon=1
+    assert len(result_df) == 3
+
+    # Check if 'future_returns' column exists
+    assert "future_returns" in result_df.columns
+
+    # Expected returns
+    # row 0: (105.0 - 100.0) / 100.0 = 0.05
+    # row 1: (102.9 - 105.0) / 105.0 = -0.02
+    # row 2: (110.0 - 102.9) / 102.9 = 0.0690...
+    expected_returns = [0.05, -0.02, (110.0 - 102.9) / 102.9]
+
+    # Assert values are close
+    np.testing.assert_allclose(result_df["future_returns"].values, expected_returns)
+
+def test_calculate_future_returns_custom_horizon():
+    """Test with custom horizon=2"""
+    df = pd.DataFrame({
+        "close": [100.0, 105.0, 110.0, 107.8]
+    })
+
+    result_df = calculate_future_returns(df.copy(), horizon=2)
+
+    # Should drop last 2 rows
+    assert len(result_df) == 2
+
+    # Expected returns for horizon 2
+    # row 0: (110.0 - 100.0) / 100.0 = 0.10
+    # row 1: (107.8 - 105.0) / 105.0 = 0.0266...
+    expected_returns = [0.10, (107.8 - 105.0) / 105.0]
+
+    np.testing.assert_allclose(result_df["future_returns"].values, expected_returns)
+
+def test_calculate_future_returns_missing_close_column():
+    """Test KeyError is raised when 'close' column is missing"""
+    df = pd.DataFrame({
+        "price": [100.0, 105.0]
+    })
+
+    with pytest.raises(KeyError):
+        calculate_future_returns(df)
+
+def test_calculate_future_returns_empty_dataframe():
+    """Test behavior with an empty dataframe"""
+    df = pd.DataFrame(columns=["close"])
+
+    result_df = calculate_future_returns(df.copy())
+
+    assert len(result_df) == 0
+    assert "future_returns" in result_df.columns
+
+def test_calculate_future_returns_all_nans():
+    """Test behavior with a dataframe containing only NaNs"""
+    df = pd.DataFrame({
+        "close": [np.nan, np.nan, np.nan]
+    })
+
+    result_df = calculate_future_returns(df.copy())
+
+    # Dropna should drop all rows
+    assert len(result_df) == 0
+    assert "future_returns" in result_df.columns
+
+def test_calculate_future_returns_horizon_larger_than_data():
+    """Test behavior when horizon is larger than dataframe size"""
+    df = pd.DataFrame({
+        "close": [100.0, 105.0]
+    })
+
+    result_df = calculate_future_returns(df.copy(), horizon=5)
+
+    # Should drop all rows
+    assert len(result_df) == 0