lumibot 4.1.3__py3-none-any.whl → 4.2.1__py3-none-any.whl

tests/test_thetadata_helper.py

@@ -13,6 +13,8 @@ import time
 from unittest.mock import patch, MagicMock
 from lumibot.entities import Asset
 from lumibot.tools import thetadata_helper
+from lumibot.backtesting import ThetaDataBacktestingPandas
+from lumibot.tools.backtest_cache import CacheMode
 
 
 @patch('lumibot.tools.thetadata_helper.update_cache')
@@ -72,11 +74,12 @@ def test_get_price_data_without_cached_data(mock_build_cache_filename, mock_get_
     # Arrange
     mock_build_cache_filename.return_value.exists.return_value = False
     mock_get_missing_dates.return_value = [datetime.datetime(2025, 9, 2)]
-    mock_get_historical_data.return_value = pd.DataFrame({
-        "datetime": pd.date_range("2023-07-01", periods=5, freq="min"),
+    raw_df = pd.DataFrame({
+        "datetime": pd.date_range("2023-07-01 09:30:00", periods=5, freq="min", tz="UTC"),
         "price": [100, 101, 102, 103, 104]
-    })
-    mock_update_df.return_value = mock_get_historical_data.return_value
+    }).set_index("datetime")
+    mock_get_historical_data.return_value = raw_df.reset_index()
+    mock_update_df.return_value = raw_df
     
     asset = Asset(asset_type="stock", symbol="AAPL")
     start = datetime.datetime(2025, 9, 2)
@@ -106,17 +109,20 @@ def test_get_price_data_partial_cache_hit(mock_build_cache_filename, mock_load_c
                                           mock_get_historical_data, mock_update_df, mock_update_cache):
     # Arrange
     cached_data = pd.DataFrame({
-        "datetime": pd.date_range("2023-07-01", periods=5, freq='min'),
-        "price": [100, 101, 102, 103, 104]
-    })
+        "datetime": pd.date_range("2023-07-01 09:30:00", periods=5, freq='min', tz="UTC"),
+        "price": [100, 101, 102, 103, 104],
+        "missing": [False] * 5,
+    }).set_index("datetime")
     mock_build_cache_filename.return_value.exists.return_value = True
     mock_load_cache.return_value = cached_data
     mock_get_missing_dates.return_value = [datetime.datetime(2025, 9, 3)]
-    mock_get_historical_data.return_value = pd.DataFrame({
-        "datetime": pd.date_range("2023-07-02", periods=5, freq='min'),
-        "price": [110, 111, 112, 113, 114]
-    })
-    updated_data = pd.concat([cached_data, mock_get_historical_data.return_value])
+    new_chunk = pd.DataFrame({
+        "datetime": pd.date_range("2023-07-02 09:30:00", periods=5, freq='min', tz="UTC"),
+        "price": [110, 111, 112, 113, 114],
+        "missing": [False] * 5,
+    }).set_index("datetime")
+    mock_get_historical_data.return_value = new_chunk.reset_index()
+    updated_data = pd.concat([cached_data, new_chunk]).sort_index()
     mock_update_df.return_value = updated_data
     
     asset = Asset(asset_type="stock", symbol="AAPL")
@@ -132,10 +138,84 @@ def test_get_price_data_partial_cache_hit(mock_build_cache_filename, mock_load_c
     assert df is not None
     assert len(df) == 10  # Combined cached and fetched data
     mock_get_historical_data.assert_called_once()
-    assert mock_update_df.return_value.equals(df)
+    pd.testing.assert_frame_equal(df, updated_data.drop(columns="missing"))
     mock_update_cache.assert_called_once()
 
 
+def test_get_price_data_daily_placeholders_prevent_refetch(monkeypatch, tmp_path):
+    from lumibot.constants import LUMIBOT_DEFAULT_PYTZ
+
+    cache_root = tmp_path / "cache_root"
+    monkeypatch.setattr(thetadata_helper, "LUMIBOT_CACHE_FOLDER", str(cache_root))
+    thetadata_helper.reset_connection_diagnostics()
+
+    asset = Asset(asset_type="stock", symbol="PLTR")
+    start = LUMIBOT_DEFAULT_PYTZ.localize(datetime.datetime(2024, 1, 1))
+    end = LUMIBOT_DEFAULT_PYTZ.localize(datetime.datetime(2024, 1, 3))
+    trading_days = [
+        datetime.date(2024, 1, 1),
+        datetime.date(2024, 1, 2),
+        datetime.date(2024, 1, 3),
+    ]
+
+    partial_df = pd.DataFrame(
+        {
+            "datetime": pd.to_datetime(["2024-01-01", "2024-01-02"], utc=True),
+            "open": [10.0, 11.0],
+            "high": [11.0, 12.0],
+            "low": [9.5, 10.5],
+            "close": [10.5, 11.5],
+            "volume": [1_000, 1_200],
+        }
+    )
+
+    progress_stub = MagicMock()
+    progress_stub.update.return_value = None
+    progress_stub.close.return_value = None
+
+    with patch("lumibot.tools.thetadata_helper.tqdm", return_value=progress_stub), \
+         patch("lumibot.tools.thetadata_helper.get_trading_dates", return_value=trading_days):
+        eod_mock = MagicMock(return_value=partial_df)
+        with patch("lumibot.tools.thetadata_helper.get_historical_eod_data", eod_mock):
+            first = thetadata_helper.get_price_data(
+                "user",
+                "pass",
+                asset,
+                start,
+                end,
+                "day",
+            )
+
+            assert eod_mock.call_count == 1
+            assert len(first) == 2
+            assert set(first.index.date) == {datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)}
+
+            cache_file = thetadata_helper.build_cache_filename(asset, "day", "ohlc")
+            loaded = thetadata_helper.load_cache(cache_file)
+            assert len(loaded) == 3
+            assert "missing" in loaded.columns
+            assert int(loaded["missing"].sum()) == 1
+            missing_dates = {idx.date() for idx, flag in loaded["missing"].items() if flag}
+            assert missing_dates == {datetime.date(2024, 1, 3)}
+
+        # Second run should reuse cache entirely
+        eod_second_mock = MagicMock(return_value=partial_df)
+        with patch("lumibot.tools.thetadata_helper.tqdm", return_value=progress_stub), \
+             patch("lumibot.tools.thetadata_helper.get_trading_dates", return_value=trading_days), \
+             patch("lumibot.tools.thetadata_helper.get_historical_eod_data", eod_second_mock):
+            second = thetadata_helper.get_price_data(
+                "user",
+                "pass",
+                asset,
+                start,
+                end,
+                "day",
+            )
+
+            assert eod_second_mock.call_count == 0
+            assert len(second) == 2
+            assert set(second.index.date) == {datetime.date(2024, 1, 1), datetime.date(2024, 1, 2)}
+
 @patch('lumibot.tools.thetadata_helper.update_cache')
 @patch('lumibot.tools.thetadata_helper.update_df')
 @patch('lumibot.tools.thetadata_helper.get_historical_data')
@@ -158,7 +238,8 @@ def test_get_price_data_empty_response(mock_build_cache_filename, mock_get_missi
     df = thetadata_helper.get_price_data("test_user", "test_password", asset, start, end, timespan, dt=dt)
 
     # Assert
-    assert df is None  # Expect None due to empty data returned
+    assert df is not None
+    assert df.empty
     mock_update_df.assert_not_called()
 
 
@@ -363,6 +444,72 @@ def test_update_cache(mocker, tmpdir, df_all, df_cached, datastyle):
     assert cache_file.exists()
 
 
+def test_get_price_data_invokes_remote_cache_manager(tmp_path, monkeypatch):
+    asset = Asset(asset_type="stock", symbol="AAPL")
+    monkeypatch.setattr(thetadata_helper, "LUMIBOT_CACHE_FOLDER", str(tmp_path))
+    cache_file = thetadata_helper.build_cache_filename(asset, "minute", "ohlc")
+    cache_file.parent.mkdir(parents=True, exist_ok=True)
+
+    df = pd.DataFrame(
+        {
+            "datetime": pd.date_range("2024-01-01 09:30:00", periods=2, freq="T", tz=pytz.UTC),
+            "open": [100.0, 101.0],
+            "high": [101.0, 102.0],
+            "low": [99.5, 100.5],
+            "close": [100.5, 101.5],
+            "volume": [1_000, 1_200],
+            "missing": [False, False],
+        }
+    )
+    df.to_parquet(cache_file, engine="pyarrow", compression="snappy", index=False)
+
+    class DummyManager:
+        def __init__(self):
+            self.ensure_calls = []
+            self.upload_calls = []
+            self.enabled = True
+            self._mode = CacheMode.S3_READWRITE
+
+        @property
+        def mode(self):
+            return self._mode
+
+        def ensure_local_file(self, local_path, payload=None, force_download=False):
+            self.ensure_calls.append((Path(local_path), payload))
+            return False
+
+        def on_local_update(self, local_path, payload=None):
+            self.upload_calls.append((Path(local_path), payload))
+            return True
+
+    dummy_manager = DummyManager()
+    monkeypatch.setattr(thetadata_helper, "get_backtest_cache", lambda: dummy_manager)
+    monkeypatch.setattr(thetadata_helper, "get_missing_dates", lambda df_all, *_args, **_kwargs: [])
+
+    start = datetime.datetime(2024, 1, 1, 9, 30, tzinfo=pytz.UTC)
+    end = datetime.datetime(2024, 1, 1, 9, 31, tzinfo=pytz.UTC)
+
+    result = thetadata_helper.get_price_data(
+        username="user",
+        password="pass",
+        asset=asset,
+        start=start,
+        end=end,
+        timespan="minute",
+        quote_asset=None,
+        dt=None,
+        datastyle="ohlc",
+        include_after_hours=True,
+        return_polars=False,
+    )
+
+    assert dummy_manager.ensure_calls, "Expected remote cache ensure call"
+    ensure_path, ensure_payload = dummy_manager.ensure_calls[0]
+    assert ensure_path == cache_file
+    assert ensure_payload["provider"] == "thetadata"
+    assert isinstance(result, pd.DataFrame)
+    assert not dummy_manager.upload_calls, "Cache hit should not trigger upload"
+
 
 @pytest.mark.parametrize(
     "df_cached, datastyle",
@@ -406,9 +553,8 @@ def test_load_data_from_cache(mocker, tmpdir, df_cached, datastyle):
     mocker.patch.object(thetadata_helper, "LUMIBOT_CACHE_FOLDER", tmpdir)
     cache_file = Path(tmpdir / "thetadata" / f"stock_SPY_1D_{datastyle}.parquet")
 
-    # No cache file
-    with pytest.raises(FileNotFoundError):
-        thetadata_helper.load_cache(cache_file)
+    # No cache file should return None (not raise)
+    assert thetadata_helper.load_cache(cache_file) is None
 
     # Cache file exists
     cache_file.parent.mkdir(parents=True, exist_ok=True)
@@ -543,7 +689,7 @@ def test_update_df_with_overlapping_data():
     assert len(df_new) == 5
     assert df_new["c"].iloc[0] == 2
     assert df_new["c"].iloc[2] == 10
-    assert df_new["c"].iloc[3] == 14 # This is the overlapping row, should keep the first value from df_all
+    assert df_new["c"].iloc[3] == 18  # Overlap prefers latest data
     assert df_new["c"].iloc[4] == 22
     # Note: The -1 minute adjustment was removed from implementation
     assert df_new.index[0] == pd.DatetimeIndex(["2025-09-02 13:30:00+00:00"])[0]
@@ -725,8 +871,12 @@ def test_get_request_error_in_json(mock_get, mock_check_connection):
 
     # Assert
     mock_get.assert_called_with(url, headers=headers, params=querystring)
-    mock_check_connection.assert_called_with(username="test_user", password="test_password")
-    assert mock_check_connection.call_count == 2
+    mock_check_connection.assert_called_with(
+        username="test_user",
+        password="test_password",
+        wait_for_connection=True,
+    )
+    assert mock_check_connection.call_count == 5
 
 
 @patch('lumibot.tools.thetadata_helper.check_connection')
@@ -744,8 +894,12 @@ def test_get_request_exception_handling(mock_get, mock_check_connection):
 
     # Assert
     mock_get.assert_called_with(url, headers=headers, params=querystring)
-    mock_check_connection.assert_called_with(username="test_user", password="test_password")
-    assert mock_check_connection.call_count == 2
+    mock_check_connection.assert_called_with(
+        username="test_user",
+        password="test_password",
+        wait_for_connection=True,
+    )
+    assert mock_check_connection.call_count == 3
 
 
 @patch('lumibot.tools.thetadata_helper.get_request')
@@ -1278,4 +1432,49 @@ class TestThetaDataChainsCaching:
 
 
 if __name__ == '__main__':
-    pytest.main()
+    pytest.main()
+
+
+def test_thetadata_no_future_minutes(monkeypatch):
+    tz = pytz.timezone('America/New_York')
+    now = tz.localize(datetime.datetime(2025, 1, 6, 9, 30))
+    frame = pd.DataFrame(
+        {
+            'datetime': [
+                tz.localize(datetime.datetime(2025, 1, 6, 9, 29)),
+                tz.localize(datetime.datetime(2025, 1, 6, 9, 31)),
+            ],
+            'open': [4330.0, 4332.0],
+            'high': [4331.0, 4333.0],
+            'low': [4329.5, 4331.5],
+            'close': [4330.5, 4332.5],
+            'volume': [1_200, 1_250],
+            'missing': [False, False],
+        }
+    )
+
+    monkeypatch.setattr(thetadata_helper, 'get_price_data', lambda *args, **kwargs: frame.copy())
+    monkeypatch.setattr(thetadata_helper, 'reset_theta_terminal_tracking', lambda: None)
+
+    data_source = ThetaDataBacktestingPandas(
+        datetime_start=now - datetime.timedelta(days=1),
+        datetime_end=now + datetime.timedelta(days=1),
+        username='user',
+        password='pass',
+        use_quote_data=False,
+    )
+    data_source._datetime = now
+
+    asset = Asset('MES', asset_type=Asset.AssetType.CONT_FUTURE)
+
+    bars = data_source.get_historical_prices(
+        asset,
+        length=1,
+        timestep='minute',
+        quote=Asset('USD', asset_type=Asset.AssetType.FOREX),
+        timeshift=datetime.timedelta(minutes=-1),
+    )
+
+    assert bars is not None
+    assert len(bars.df) == 1
+    assert bars.df.index[-1].tz_convert(tz) <= now

tests/test_thetadata_pandas_verification.py

@@ -0,0 +1,186 @@
+"""
+Verification test for ThetaData pandas implementation.
+
+This test verifies that the pandas implementation:
+1. Works correctly with caching (cold→warm produces 0 requests)
+2. Produces consistent results between cold and warm runs
+3. Returns correct data for the WeeklyMomentumOptionsStrategy symbols
+
+This establishes the baseline before cloning to polars.
+"""
+
+import os
+import shutil
+import json
+from datetime import datetime
+from pathlib import Path
+import pytest
+
+from lumibot.backtesting import ThetaDataBacktestingPandas
+from lumibot.strategies import Strategy
+from lumibot.entities import Asset
+from lumibot.credentials import THETADATA_CONFIG
+
+
+def get_cache_dir():
+    """Get the ThetaData cache directory."""
+    return Path.home() / "Library" / "Caches" / "lumibot" / "1.0" / "thetadata"
+
+
+def clear_cache():
+    """Clear all ThetaData cache files."""
+    cache_dir = get_cache_dir()
+    if cache_dir.exists():
+        print(f"Clearing cache at {cache_dir}")
+        shutil.rmtree(cache_dir)
+        cache_dir.mkdir(parents=True, exist_ok=True)
+    else:
+        cache_dir.mkdir(parents=True, exist_ok=True)
+    print("Cache cleared")
+
+
+def count_cache_files():
+    """Count the number of cache files."""
+    cache_dir = get_cache_dir()
+    if not cache_dir.exists():
+        return 0
+    return len(list(cache_dir.glob("*.parquet")))
+
+
+class WeeklyMomentumOptionsStrategy(Strategy):
+    """Simplified version of WeeklyMomentumOptionsStrategy for testing."""
+
+    def initialize(self):
+        self.sleeptime = "1D"
+        self.data_fetches = []
+        self.symbols = ["SPY", "QQQ", "IWM"]
+
+    def on_trading_iteration(self):
+        # Fetch historical data for each symbol
+        for symbol in self.symbols:
+            asset = Asset(symbol, asset_type=Asset.AssetType.STOCK)
+
+            # Get 5 days of daily data
+            daily_bars = self.get_historical_prices(asset, length=5, timestep="day")
+            if daily_bars and hasattr(daily_bars, 'df'):
+                self.data_fetches.append({
+                    "symbol": symbol,
+                    "timestep": "day",
+                    "length": 5,
+                    "rows": len(daily_bars.df)
+                })
+
+            # Get 10 minutes of minute data
+            minute_bars = self.get_historical_prices(asset, length=10, timestep="minute")
+            if minute_bars and hasattr(minute_bars, 'df'):
+                self.data_fetches.append({
+                    "symbol": symbol,
+                    "timestep": "minute",
+                    "length": 10,
+                    "rows": len(minute_bars.df)
+                })
+
+
+def run_backtest(run_type):
+    """Run a backtest and return the strategy data_fetches."""
+    print(f"\n{'='*60}")
+    print(f"Running {run_type.upper()} backtest with pandas")
+    print(f"{'='*60}")
+
+    cache_before = count_cache_files()
+    print(f"Cache files before: {cache_before}")
+
+    # Run backtest using Strategy.run_backtest() class method to get both results and strategy
+    results, strategy_instance = WeeklyMomentumOptionsStrategy.run_backtest(
+        ThetaDataBacktestingPandas,
+        backtesting_start=datetime(2025, 3, 1),
+        backtesting_end=datetime(2025, 3, 14),
+        budget=100000,
+        show_plot=False,
+        show_tearsheet=False,
+        save_tearsheet=False,
+        show_indicators=False,
+        quiet_logs=True,
+        show_progress_bar=False,
+    )
+
+    cache_after = count_cache_files()
+    print(f"Cache files after: {cache_after}")
+    print(f"New cache files created: {cache_after - cache_before}")
+
+    # Get portfolio value from strategy instance
+    portfolio_value = strategy_instance.portfolio_value
+
+    # Get data fetches count
+    data_fetches = len(strategy_instance.data_fetches) if hasattr(strategy_instance, 'data_fetches') else 0
+
+    print(f"Portfolio value: ${portfolio_value:,.2f}")
+    print(f"Data fetches: {data_fetches}")
+    print(f"Results: {results}")
+
+    return {
+        "run_type": run_type,
+        "portfolio_value": portfolio_value,
+        "data_fetches": data_fetches,
+        "cache_before": cache_before,
+        "cache_after": cache_after,
+        "new_cache_files": cache_after - cache_before,
+        "fetch_details": strategy_instance.data_fetches if hasattr(strategy_instance, 'data_fetches') else [],
+        "results": results
+    }
+
+
+@pytest.mark.apitest
+@pytest.mark.skipif(
+    not THETADATA_CONFIG.get("THETADATA_USERNAME") or not THETADATA_CONFIG.get("THETADATA_PASSWORD"),
+    reason="ThetaData credentials not configured - skipping API test"
+)
+def test_pandas_cold_warm():
+    """Test that pandas implementation works correctly with caching."""
+
+    # Clear cache and run cold
+    clear_cache()
+    cold_results = run_backtest("cold")
+
+    # Run warm (cache should be used)
+    warm_results = run_backtest("warm")
+
+    # Verify results
+    print(f"\n{'='*60}")
+    print("VERIFICATION RESULTS")
+    print(f"{'='*60}")
+
+    # Check 1: Cold run should create cache files
+    assert cold_results["new_cache_files"] > 0, "Cold run should create cache files"
+    print(f"✓ Cold run created {cold_results['new_cache_files']} cache files")
+
+    # Check 2: Warm run should not create new cache files
+    assert warm_results["new_cache_files"] == 0, "Warm run should not create new cache files"
+    print(f"✓ Warm run created {warm_results['new_cache_files']} new cache files (expected 0)")
+
+    # Check 3: Portfolio values should match
+    pv_diff = abs(cold_results["portfolio_value"] - warm_results["portfolio_value"])
+    assert pv_diff < 0.01, f"Portfolio values should match (diff: ${pv_diff:,.2f})"
+    print(f"✓ Portfolio values match: ${cold_results['portfolio_value']:,.2f}")
+
+    # Check 4: Data fetches should match
+    assert cold_results["data_fetches"] == warm_results["data_fetches"], "Data fetches should match"
+    print(f"✓ Data fetches match: {cold_results['data_fetches']}")
+
+    # Save results for reference
+    results_path = Path("/Users/robertgrzesik/Documents/Development/lumivest_bot_server/strategies/lumibot/logs/pandas_verification_results.json")
+    results_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(results_path, 'w') as f:
+        json.dump({
+            "cold": cold_results,
+            "warm": warm_results
+        }, f, indent=2, default=str)
+
+    print(f"\n✓ Results saved to {results_path}")
+    print("\n✅ ALL CHECKS PASSED - Pandas implementation is working correctly")
+
+    return cold_results, warm_results
+
+
+if __name__ == "__main__":
+    test_pandas_cold_warm()