lumibot 4.1.2__py3-none-any.whl → 4.2.0__py3-none-any.whl

tests/test_databento_helper.py

@@ -49,7 +49,12 @@ class TestDataBentoHelper(unittest.TestCase):
         self.assertRegex(result, r"ES[FGHJKMNQUVXZ]\d", "Should auto-resolve to contract format like ESZ5")
         
         # Test specific contract with expiration (March 2025 = H25)
-        result = databento_helper._format_futures_symbol_for_databento(self.test_asset_future)
+        specific_future = Asset(
+            symbol="ES",
+            asset_type="future",
+            expiration=datetime(2025, 3, 15).date()
+        )
+        result = databento_helper._format_futures_symbol_for_databento(specific_future)
         self.assertEqual(result, "ESH25")  # March 2025 = H25
         
         # Test another month (December 2024 = Z24)

tests/test_databento_live.py

@@ -27,13 +27,13 @@ load_dotenv()
 def test_symbol_resolution():
     """Test that symbols are properly resolved to contract codes"""
     from lumibot.entities import Asset
-    from lumibot.data_sources.databento_data_polars_live import DataBentoDataPolarsLive
+    from lumibot.data_sources.databento_data_polars import DataBentoDataPolars
     
     print("\n" + "="*60)
     print("TEST 1: Symbol Resolution")
     print("="*60)
     
-    data_source = DataBentoDataPolarsLive(
+    data_source = DataBentoDataPolars(
         api_key=os.getenv('DATABENTO_API_KEY'),
         has_paid_subscription=True,
         enable_live_stream=False  # Don't need streaming for this test
@@ -97,14 +97,14 @@ def test_live_api_connection():
 def test_minute_bar_aggregation():
     """Test minute bar aggregation with <1 minute lag"""
     from lumibot.entities import Asset
-    from lumibot.data_sources.databento_data_polars_live import DataBentoDataPolarsLive
+    from lumibot.data_sources.databento_data_polars import DataBentoDataPolars
     
     print("\n" + "="*60)
     print("TEST 3: Minute Bar Aggregation & Latency")
     print("="*60)
     
     # Initialize with Live API
-    data_source = DataBentoDataPolarsLive(
+    data_source = DataBentoDataPolars(
         api_key=os.getenv('DATABENTO_API_KEY'),
         has_paid_subscription=True,
         enable_live_stream=True
@@ -175,13 +175,13 @@ def test_minute_bar_aggregation():
 )
 def test_api_routing():
     """Test that correct API is used based on time range"""
-    from lumibot.data_sources.databento_data_polars_live import DataBentoDataPolarsLive
+    from lumibot.data_sources.databento_data_polars import DataBentoDataPolars
     
     print("\n" + "="*60)
     print("TEST 4: API Routing (Live vs Historical)")
     print("="*60)
     
-    data_source = DataBentoDataPolarsLive(
+    data_source = DataBentoDataPolars(
         api_key=os.getenv('DATABENTO_API_KEY'),
         has_paid_subscription=True,
         enable_live_stream=True
@@ -219,13 +219,13 @@ def test_api_routing():
 def test_long_time_periods():
     """Test different time periods including long periods (500+ bars)"""
     from lumibot.entities import Asset
-    from lumibot.data_sources.databento_data_polars_live import DataBentoDataPolarsLive
+    from lumibot.data_sources.databento_data_polars import DataBentoDataPolars
     
     print("\n" + "="*60)
     print("TEST 5: Long Time Period Handling (500+ bars)")
     print("="*60)
     
-    data_source = DataBentoDataPolarsLive(
+    data_source = DataBentoDataPolars(
         api_key=os.getenv('DATABENTO_API_KEY'),
         has_paid_subscription=True,
         enable_live_stream=True
@@ -324,14 +324,14 @@ def test_long_time_periods():
 def test_continuous_latency_monitoring():
     """Run continuous tests to verify consistent <1 minute lag"""
     from lumibot.entities import Asset
-    from lumibot.data_sources.databento_data_polars_live import DataBentoDataPolarsLive
+    from lumibot.data_sources.databento_data_polars import DataBentoDataPolars
     
     print("\n" + "="*60)
     print("TEST 6: Continuous Latency Monitoring")
     print("="*60)
     print("Running 5 consecutive tests to verify consistent low latency...")
     
-    data_source = DataBentoDataPolarsLive(
+    data_source = DataBentoDataPolars(
         api_key=os.getenv('DATABENTO_API_KEY'),
         has_paid_subscription=True,
         enable_live_stream=True

tests/test_futures_roll.py

@@ -0,0 +1,38 @@
+import datetime
+
+import pytz
+
+from lumibot.entities import Asset
+from lumibot.tools import futures_roll
+
+NY = pytz.timezone("America/New_York")
+
+
+def _dt(year: int, month: int, day: int) -> datetime.datetime:
+    return NY.localize(datetime.datetime(year, month, day))
+
+
+def test_equity_index_roll_eight_business_days_before_expiry():
+    asset = Asset("MES", asset_type=Asset.AssetType.CONT_FUTURE)
+
+    year, month = futures_roll.determine_contract_year_month(asset.symbol, _dt(2025, 9, 8))
+    assert (year, month) == (2025, 9)
+
+    year, month = futures_roll.determine_contract_year_month(asset.symbol, _dt(2025, 9, 10))
+    assert (year, month) == (2025, 12)
+
+
+def test_fallback_mid_month_preserved_for_unknown_symbols():
+    asset = Asset("XYZ", asset_type=Asset.AssetType.CONT_FUTURE)
+
+    year, month = futures_roll.determine_contract_year_month(asset.symbol, _dt(2025, 3, 16))
+    assert (year, month) == (2025, 6)
+
+
+def test_resolve_symbols_for_range_produces_sequential_contracts():
+    asset = Asset("MES", asset_type=Asset.AssetType.CONT_FUTURE)
+    start = _dt(2025, 8, 1)
+    end = _dt(2025, 12, 31)
+
+    symbols = futures_roll.resolve_symbols_for_range(asset, start, end, year_digits=1)
+    assert symbols == ["MESU5", "MESZ5", "MESH6"], symbols

tests/test_indicator_subplots.py

@@ -3,10 +3,12 @@ from datetime import datetime as DateTime
 from unittest.mock import MagicMock
 
 import pandas as pd
+import plotly.graph_objects as go
 
 from lumibot.backtesting import PandasDataBacktesting
 from lumibot.strategies.strategy import Strategy
 from lumibot.entities import Asset
+from lumibot.tools.indicators import _build_trade_marker_tooltip, plot_returns
 
 from tests.fixtures import pandas_data_fixture
 
@@ -179,6 +181,105 @@ class TestIndicators:
         logger.info(f"Result: {result}")
         assert result is not None
 
+
+def _make_trade_row_for_tooltip(status, trade_cost=pd.NA):
+    return pd.Series(
+        {
+            "status": status,
+            "filled_quantity": 10,
+            "price": 2.5,
+            "asset.multiplier": 100,
+            "trade_cost": trade_cost,
+            "symbol": "WDC",
+            "asset.asset_type": "option",
+            "asset.right": "CALL",
+            "asset.strike": 86,
+            "asset.expiration": "2025-09-19",
+            "type": "limit",
+        }
+    )
+
+
+def test_cash_settled_tooltip_generated_without_trade_cost():
+    tooltip = _build_trade_marker_tooltip(_make_trade_row_for_tooltip("cash_settled", trade_cost=pd.NA))
+    assert tooltip is not None
+    assert "cash_settled" in tooltip
+
+
+def test_non_terminal_status_filtered_out():
+    assert _build_trade_marker_tooltip(_make_trade_row_for_tooltip("new", trade_cost=pd.NA)) is None
+
+
+def test_plot_returns_preserves_cash_settled_status(tmp_path, monkeypatch):
+    plot_path = tmp_path / "plot.html"
+
+    def _fake_write_html(self, file, auto_open=True, **kwargs):
+        # Prevent plotly from opening a browser during the test
+        return file
+
+    monkeypatch.setattr(go.Figure, "write_html", _fake_write_html, raising=False)
+
+    idx = pd.to_datetime(
+        ["2025-09-04 00:00:00-04:00", "2025-09-20 00:00:00-04:00"]
+    ).tz_convert("UTC")
+
+    strategy_df = pd.DataFrame(
+        {
+            "return": [0.0, 0.0],
+            "cash": [100000, 120000],
+            "positions": [
+                [{"asset": "WDC", "quantity": 25}],
+                [],
+            ],
+        },
+        index=idx,
+    )
+
+    benchmark_df = pd.DataFrame(
+        {
+            "return": [0.0, 0.0],
+            "open": [1.0, 1.0],
+            "high": [1.0, 1.0],
+            "low": [1.0, 1.0],
+            "close": [1.0, 1.0],
+        },
+        index=idx,
+    )
+
+    trades_df = pd.DataFrame(
+        [
+            {
+                "time": "2025-09-20 00:00:00-04:00",
+                "side": "sell",
+                "status": "cash_settled",
+                "filled_quantity": 25,
+                "symbol": "WDC",
+                "asset.asset_type": "option",
+                "asset.right": "CALL",
+                "asset.strike": 86,
+                "asset.expiration": "2025-09-19",
+                "price": 20.86,
+                "type": "cash_settled",
+                "asset.multiplier": 100,
+                "trade_cost": pd.NA,
+            }
+        ]
+    )
+
+    plot_returns(
+        strategy_df,
+        "Strategy",
+        benchmark_df,
+        "Benchmark",
+        plot_file_html=str(plot_path),
+        trades_df=trades_df,
+        show_plot=True,
+        initial_budget=1,
+    )
+
+    trades_csv = pd.read_csv(plot_path.with_suffix(".csv"))
+    assert "cash_settled" in trades_csv["status"].tolist()
+
     def test_named_lines(self, pandas_data_fixture):
         """Test the named lines"""
         strategy_name = "TestIndicatorStrategy"

tests/test_market_infinite_loop_bug.py

@@ -16,11 +16,16 @@ STATUS: ✅ FIXED - ES futures now complete normally (1 restart vs infinite)
 
 import unittest
 from unittest.mock import patch
-from datetime import datetime
+from datetime import datetime, timedelta
 
+import pandas as pd
+
+from lumibot.credentials import DATABENTO_CONFIG
 from lumibot.strategies import Strategy
-from lumibot.entities import Asset, TradingFee
-from lumibot.backtesting import DataBentoDataBacktesting
+from lumibot.entities import Asset, TradingFee, Bars, Order
+from lumibot.backtesting import BacktestingBroker, DataBentoDataBacktesting
+
+DATABENTO_API_KEY = DATABENTO_CONFIG.get("API_KEY")
 
 
 class ESFuturesTestStrategy(Strategy):
@@ -39,10 +44,13 @@ class TestESFuturesHangBug(unittest.TestCase):
     """Test that ES futures strategies no longer hang/restart infinitely"""
     
     def setUp(self):
+        if not DATABENTO_API_KEY or DATABENTO_API_KEY == "<your key here>":
+            self.skipTest("DataBento API key required for DataBento backtesting tests")
         self.backtesting_params = {
             'datasource_class': DataBentoDataBacktesting,
             'backtesting_start': datetime(2025, 6, 5),
             'backtesting_end': datetime(2025, 6, 6),
+            'api_key': DATABENTO_API_KEY,
             'show_plot': False,
             'show_tearsheet': False,
             'show_indicators': False,
@@ -299,3 +307,69 @@ class TestESFuturesHangBug(unittest.TestCase):
 if __name__ == '__main__':
     print("🧪 Testing ES Futures hang bug fix...")
     unittest.main(verbosity=2)
+
+
+def test_broker_timeshift_guard():
+    captured = []
+
+    class StubDataSource:
+        SOURCE = "DATABENTO_POLARS"
+        IS_BACKTESTING_DATA_SOURCE = True
+
+        def __init__(self):
+            self._datetime = datetime(2025, 6, 5, 14, 30)
+
+        def get_historical_prices(self, asset, length, quote=None, timeshift=None, **kwargs):
+            captured.append(timeshift)
+            index = pd.DatetimeIndex([self._datetime - timedelta(minutes=1)])
+            frame = pd.DataFrame(
+                {
+                    'open': [4300.0],
+                    'high': [4301.0],
+                    'low': [4299.5],
+                    'close': [4300.5],
+                    'volume': [1500],
+                },
+                index=index,
+            )
+            target_asset = asset[0] if isinstance(asset, tuple) else asset
+            return Bars(frame, self.SOURCE, target_asset, raw=frame)
+
+        def get_datetime(self):
+            return self._datetime
+
+    broker = BacktestingBroker(data_source=StubDataSource())
+    broker._datetime = broker.data_source.get_datetime()
+
+    order = Order(
+        strategy="stub",
+        asset=Asset("MES", asset_type=Asset.AssetType.CONT_FUTURE),
+        quantity=1,
+        side=Order.OrderSide.BUY,
+    )
+    order.order_type = Order.OrderType.MARKET
+    order.quote = Asset("USD", asset_type=Asset.AssetType.FOREX)
+    broker._new_orders.append(order)
+
+    class StubStrategy:
+        name = "stub"
+        buy_trading_fees = []
+        sell_trading_fees = []
+        timestep = 'minute'
+        bars_lookback = 1
+
+        def __init__(self, broker):
+            self.broker = broker
+            self.cash = 100000.0
+            self.quote_asset = Asset('USD', asset_type=Asset.AssetType.FOREX)
+
+        def log_message(self, *args, **kwargs):
+            return None
+
+        def _set_cash_position(self, value):
+            self.cash = value
+
+    broker.process_pending_orders(strategy=StubStrategy(broker))
+
+    assert captured, "BacktestingBroker did not request historical data"
+    assert captured[0] == timedelta(minutes=-2)

tests/test_polars_resample.py

@@ -0,0 +1,67 @@
+from datetime import datetime, timedelta
+
+import polars as pl
+import pytest
+
+from lumibot.tools.polars_utils import PolarsResampleError, resample_polars_ohlc
+
+
+def _minute_frame(rows: int = 10) -> pl.DataFrame:
+    start = datetime(2024, 1, 1, 9, 30)
+    datetimes = [start + timedelta(minutes=i) for i in range(rows)]
+    return pl.DataFrame(
+        {
+            "datetime": datetimes,
+            "open": [100 + i for i in range(rows)],
+            "high": [101 + i for i in range(rows)],
+            "low": [99 + i for i in range(rows)],
+            "close": [100.5 + i for i in range(rows)],
+            "volume": [1_000 + 10 * i for i in range(rows)],
+            "signal": [i for i in range(rows)],
+        }
+    )
+
+
+def test_resample_polars_minute_to_5min():
+    df = _minute_frame(10)
+    resampled = resample_polars_ohlc(df, multiplier=5, base_unit="minute", length=2)
+
+    assert resampled.height == 2
+    # First bucket should cover rows 0-4
+    first = resampled.row(0, named=True)
+    assert first["open"] == 100
+    assert first["high"] == 105
+    assert first["low"] == 99
+    assert first["close"] == pytest.approx(104.5)
+    assert first["volume"] == sum(1_000 + 10 * i for i in range(5))
+    # signal column keeps last observation in the bucket
+    assert first["signal"] == 4
+
+    # Tail limiting should keep last 2 buckets only
+    assert resampled.row(1, named=True)["signal"] == 9
+
+
+def test_resample_polars_day_bucket():
+    start = datetime(2024, 1, 1, 0, 0)
+    datetimes = [start + timedelta(hours=i) for i in range(48)]
+    df = pl.DataFrame(
+        {
+            "datetime": datetimes,
+            "open": [10 + i for i in range(48)],
+            "high": [12 + i for i in range(48)],
+            "low": [8 + i for i in range(48)],
+            "close": [11 + i for i in range(48)],
+            "volume": [100 + i for i in range(48)],
+        }
+    )
+
+    resampled = resample_polars_ohlc(df, multiplier=1, base_unit="day", length=None)
+    assert resampled.height == 2
+    assert resampled["open"][0] == 10
+    assert resampled["close"][1] == 11 + 47
+
+
+def test_resample_polars_invalid_unit():
+    df = _minute_frame(2)
+    with pytest.raises(PolarsResampleError):
+        resample_polars_ohlc(df, multiplier=1, base_unit="hour")

tests/test_polygon_helper.py

@@ -7,6 +7,7 @@ import pytest
 import pytz
 
 from lumibot.entities import Asset
+from lumibot.backtesting import PolygonDataBacktesting
 from lumibot.entities.chains import normalize_option_chains
 from lumibot.tools import polygon_helper as ph
 
@@ -131,6 +132,8 @@ class TestPolygonHelpers:
         assert datetime.date(2023, 7, 10) in trading_dates
 
     def test_get_polygon_symbol(self, mocker):
+
+
         polygon_client = mocker.MagicMock()
 
         # ------- Unsupported Asset Type
@@ -568,3 +571,46 @@ class TestPolygonPriceData:
         # Should return identical data once normalized
         assert normalize_option_chains(result_second) == normalized_first
         assert mock_polyclient.list_options_contracts.call_count == 0
+
+    def test_polygon_no_future_bars_before_open(self, monkeypatch):
+        tz = pytz.timezone('America/New_York')
+        now = tz.localize(datetime.datetime(2023, 11, 1, 9, 30))
+        frame = pd.DataFrame(
+            {
+                'open': [377.0, 378.5],
+                'high': [377.5, 379.0],
+                'low': [376.8, 378.2],
+                'close': [377.2, 378.9],
+                'volume': [10_000, 10_500],
+            },
+            index=pd.DatetimeIndex([
+                tz.localize(datetime.datetime(2023, 11, 1, 9, 29)),
+                tz.localize(datetime.datetime(2023, 11, 1, 9, 31)),
+            ]),
+        )
+
+        monkeypatch.setattr(
+            'lumibot.backtesting.polygon_backtesting.polygon_helper.get_price_data_from_polygon',
+            lambda *args, **kwargs: frame,
+        )
+
+        data_source = PolygonDataBacktesting(
+            datetime_start=now - datetime.timedelta(days=1),
+            datetime_end=now + datetime.timedelta(days=1),
+            api_key='dummy',
+        )
+        data_source._datetime = now
+        asset = Asset('SPY')
+        quote = Asset('USD', 'forex')
+
+        bars = data_source.get_historical_prices(
+            asset,
+            length=1,
+            timestep='minute',
+            quote=quote,
+            timeshift=datetime.timedelta(minutes=-1),
+        )
+
+        assert bars.df.index[-1] <= now
+
+

tests/test_thetadata_backwards_compat.py

@@ -0,0 +1,97 @@
+"""
+ThetaData pandas compatibility tests.
+
+Ensures the helper returns pandas DataFrames by default and raises a clear
+error when callers request polars output (which is intentionally unsupported
+in this branch).
+"""
+
+from datetime import datetime, timezone
+import pandas as pd
+import pytest
+
+from lumibot.entities import Asset
+from lumibot.tools import thetadata_helper
+
+
+def _mock_cache_frame(start: datetime, rows: int = 8) -> pd.DataFrame:
+    index = pd.date_range(start=start, periods=rows, freq="1min", tz="UTC")
+    df = pd.DataFrame(
+        {
+            "open": [200 + i for i in range(rows)],
+            "high": [200.5 + i for i in range(rows)],
+            "low": [199.5 + i for i in range(rows)],
+            "close": [200.25 + i for i in range(rows)],
+            "volume": [10_000 + 50 * i for i in range(rows)],
+            "missing": [False] * rows,
+        },
+        index=index,
+    )
+    return df
+
+
+def test_get_price_data_returns_pandas_when_cache_hit(monkeypatch, tmp_path):
+    """Cache path with no missing intervals should return pandas DataFrame."""
+    cache_file = tmp_path / "spy.minute.ohlc.parquet"
+    cache_file.write_text("placeholder")
+
+    mock_df = _mock_cache_frame(datetime(2025, 1, 1, tzinfo=timezone.utc))
+
+    monkeypatch.setattr(
+        thetadata_helper,
+        "build_cache_filename",
+        lambda *args, **kwargs: cache_file,
+    )
+    monkeypatch.setattr(thetadata_helper, "load_cache", lambda _: mock_df)
+    monkeypatch.setattr(thetadata_helper, "get_missing_dates", lambda *args, **kwargs: [])
+    monkeypatch.setattr(thetadata_helper, "update_cache", lambda *args, **kwargs: None)
+
+    asset = Asset("SPY", asset_type=Asset.AssetType.STOCK)
+
+    result = thetadata_helper.get_price_data(
+        username="demo",
+        password="demo",
+        asset=asset,
+        start=datetime(2025, 1, 1, tzinfo=timezone.utc),
+        end=datetime(2025, 1, 2, tzinfo=timezone.utc),
+        timespan="minute",
+        datastyle="ohlc",
+        include_after_hours=True,
+    )
+
+    assert isinstance(result, pd.DataFrame)
+    assert not result.empty
+    assert "missing" not in result.columns
+
+
+def test_get_price_data_polars_request_rejected(monkeypatch, tmp_path):
+    """Requesting return_polars=True should raise a clear ValueError."""
+    cache_file = tmp_path / "spy.minute.ohlc.parquet"
+    cache_file.write_text("placeholder")
+
+    mock_df = _mock_cache_frame(datetime(2025, 1, 1, tzinfo=timezone.utc))
+
+    monkeypatch.setattr(
+        thetadata_helper,
+        "build_cache_filename",
+        lambda *args, **kwargs: cache_file,
+    )
+    monkeypatch.setattr(thetadata_helper, "load_cache", lambda _: mock_df)
+    monkeypatch.setattr(thetadata_helper, "get_missing_dates", lambda *args, **kwargs: [])
+
+    asset = Asset("SPY", asset_type=Asset.AssetType.STOCK)
+
+    with pytest.raises(ValueError) as excinfo:
+        thetadata_helper.get_price_data(
+            username="demo",
+            password="demo",
+            asset=asset,
+            start=datetime(2025, 1, 1, tzinfo=timezone.utc),
+            end=datetime(2025, 1, 2, tzinfo=timezone.utc),
+            timespan="minute",
+            datastyle="ohlc",
+            include_after_hours=True,
+            return_polars=True,
+        )
+
+    assert "polars output" in str(excinfo.value).lower()