Implement test cases for merge_data_with_backtest_portfolio and move the main function to backtest_signals.py.

Author: niklasniedermeier

src/backtest_bay/backtest/backtest_signals.py

@@ -236,3 +236,16 @@ def _validate_price_col(data, price_col):
     if price_col not in data.columns:
         error_msg = f"data must contain a '{price_col}' column."
         raise ValueError(error_msg)
+
+
+def merge_data_with_backtest_portfolio(data, portfolio):
+    """Merge downloaded data with backtested portfolio using the index.
+
+    Args:
+        data (pd.DataFrame): DataFrame with downloaded data.
+        portfolio (pd.DataFrame): DataFrame to be merged with data using the index.
+
+    Returns:
+        pd.DataFrame: Merged DataFrame.
+    """
+    return data.merge(portfolio, how="left", left_index=True, right_index=True)

src/backtest_bay/backtest/task_backtest.py

@@ -1,7 +1,10 @@
 import pandas as pd
 import pytask
 
-from backtest_bay.backtest.backtest_signals import backtest_signals
+from backtest_bay.backtest.backtest_signals import (
+    backtest_signals,
+    merge_data_with_backtest_portfolio,
+)
 from backtest_bay.backtest.generate_signals import generate_signals
 from backtest_bay.config import BLD, INITIAL_CASH, PARAMS, SRC, TAC, TRADE_PCT
 
@@ -38,21 +41,8 @@ def task_backtest(
             trade_pct=TRADE_PCT,
         )
 
-        merged_portfolio = _merge_stock_data_with_portfolio(
+        merged_portfolio = merge_data_with_backtest_portfolio(
             stock_data, backtested_portfolio
         )
 
         merged_portfolio.to_pickle(produces)
-
-
-def _merge_stock_data_with_portfolio(data, portfolio):
-    """Merge data with portfolio using the index.
-
-    Args:
-        data (pd.DataFrame): DataFrame with downloaded data.
-        portfolio (pd.DataFrame): DataFrame to be merged with data using the index.
-
-    Returns:
-        pd.DataFrame: Merged DataFrame.
-    """
-    return data.merge(portfolio, how="left", left_index=True, right_index=True)

tests/backtest/test_backtest_signals.py

@@ -13,6 +13,7 @@
     _validate_tac,
     _validate_trade_pct,
     backtest_signals,
+    merge_data_with_backtest_portfolio,
 )
 
 
@@ -225,3 +226,29 @@ def test_validate_trade_pct_invalid_input(trade_pct, expected_error):
 def test_validate_data_valid_input(data, price_col):
     """Test valid data input for _validate_data."""
     _validate_price_col(data, price_col)
+
+
+# Tests for merge_data_with_backtest_portfolio
+@pytest.mark.parametrize(
+    ("data", "portfolio", "expected"),
+    [
+        (
+            pd.DataFrame({"price": [10, 6]}, index=["2023-01-01", "2023-01-02"]),
+            pd.DataFrame({"signal": [1, 0]}, index=["2023-01-01", "2023-01-02"]),
+            pd.DataFrame(
+                {"price": [10, 6], "signal": [1, 0]}, index=["2023-01-01", "2023-01-02"]
+            ),
+        ),
+        (
+            pd.DataFrame({"price": [5, 44]}, index=["2023-05-01", "2023-05-08"]),
+            pd.DataFrame({"signal": [1, 1]}, index=["2023-05-01", "2023-05-08"]),
+            pd.DataFrame(
+                {"price": [5, 44], "signal": [1, 1]}, index=["2023-05-01", "2023-05-08"]
+            ),
+        ),
+    ],
+)
+def test_merge_data_with_backtest_portfolio(data, portfolio, expected):
+    "Test correct merge for merge_data_with_backtest_portfolio."
+    result = merge_data_with_backtest_portfolio(data, portfolio)
+    result.equals(expected)