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

tests/test_databento_data.py

@@ -1,493 +1,250 @@
+import os
 import unittest
-from unittest.mock import Mock, patch, MagicMock
-from datetime import datetime, timedelta
+from datetime import datetime, timedelta, timezone
+from unittest.mock import MagicMock, patch
+
 import pandas as pd
+import polars as pl
 
-from lumibot.data_sources.databento_data import DataBentoData
+from lumibot.data_sources import DataBentoData
 from lumibot.entities import Asset, Bars
 
 
 class TestDataBentoData(unittest.TestCase):
-    """Test cases for DataBentoData data source"""
+    """Unit tests for the canonical DataBento data source (Polars-backed)."""
 
     def setUp(self):
-        """Set up test fixtures"""
         self.api_key = "test_api_key"
-        self.start_date = datetime(2025, 1, 1)
-        self.end_date = datetime(2025, 1, 31)
-        
-        self.test_asset = Asset(
+        patchers = [
+            patch("lumibot.tools.databento_helper.DATABENTO_AVAILABLE", True),
+            patch("lumibot.tools.databento_helper_polars.DATABENTO_AVAILABLE", True),
+            patch("lumibot.tools.databento_helper_polars.DataBentoClient", MagicMock()),
+        patch("lumibot.tools.databento_helper_polars._fetch_and_update_futures_multiplier", lambda *args, **kwargs: None),
+        ]
+        for patcher in patchers:
+            patched = patcher.start()
+            self.addCleanup(patcher.stop)
+
+        import importlib
+
+        polars_module = importlib.import_module("lumibot.data_sources.databento_data_polars")
+        patcher_db = patch.object(polars_module, "db", MagicMock())
+        patcher_db.start()
+        self.addCleanup(patcher_db.stop)
+
+        self.future_asset = Asset(
             symbol="ES",
-            asset_type="future",
-            expiration=datetime(2025, 3, 15).date()
-        )
+            asset_type=Asset.AssetType.FUTURE,
+            expiration=datetime(2025, 3, 15).date(),
+        )
+        self.cont_future_asset = Asset(
+            symbol="MES",
+            asset_type=Asset.AssetType.CONT_FUTURE,
+        )
+        self.equity_asset = Asset("AAPL", asset_type=Asset.AssetType.STOCK)
+        # Disable live streaming threads for unit-speed tests
+        self.datasource_kwargs = {"api_key": self.api_key, "enable_live_stream": False}
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _polars_ohlcv(rows: int = 3) -> pl.DataFrame:
+        base_time = datetime(2025, 1, 1, 9, 30, tzinfo=timezone.utc)
+        minutes = [base_time + timedelta(minutes=i) for i in range(rows)]
+        return pl.DataFrame(
+            {
+                "datetime": minutes,
+                "open": [100.0 + i for i in range(rows)],
+                "high": [101.0 + i for i in range(rows)],
+                "low": [99.0 + 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)],
+            }
+        )
+
+    def _bars(self, rows: int = 2) -> Bars:
+        df = self._polars_ohlcv(rows)
+        return Bars(df=df, source="DATABENTO", asset=self.future_asset)
+
+    # ------------------------------------------------------------------
+    # Initialization
+    # ------------------------------------------------------------------
+    def test_initialization_sets_core_attributes(self):
+        data_source = DataBentoData(**self.datasource_kwargs)
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_initialization_success(self):
-        """Test successful initialization"""
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        self.assertEqual(data_source.name, "databento")
-        self.assertEqual(data_source.SOURCE, "DATABENTO")
         self.assertEqual(data_source._api_key, self.api_key)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', False)
-    def test_initialization_databento_unavailable(self):
-        """Test initialization when DataBento is unavailable"""
-        with self.assertRaises(ImportError):
-            DataBentoData(
-                api_key=self.api_key,
-                datetime_start=self.start_date,
-                datetime_end=self.end_date
+        self.assertEqual(data_source.SOURCE, "DATABENTO")
+        # Live streaming disabled for tests should be reflected on the instance
+        self.assertFalse(data_source.enable_live_stream)
+        # Name comes from DataSource base class
+        self.assertEqual(data_source.name, "data_source")
+
+    # ------------------------------------------------------------------
+    # Historical data
+    # ------------------------------------------------------------------
+    def test_get_historical_prices_returns_bars(self):
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            return_value=self._polars_ohlcv(3),
+        ) as mock_get_data:
+            data_source = DataBentoData(**self.datasource_kwargs)
+            bars = data_source.get_historical_prices(
+                asset=self.future_asset,
+                length=3,
+                timestep="minute",
             )
 
-    def test_initialization_default_dates(self):
-        """Test initialization with default dates"""
-        with patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True):
-            data_source = DataBentoData(api_key=self.api_key)
-            
-            # Should have set API key and other attributes
-            self.assertIsNotNone(data_source._api_key)
-            self.assertEqual(data_source._api_key, self.api_key)
-            self.assertIsNotNone(data_source.name)
-            self.assertEqual(data_source.name, "databento")
-            self.assertFalse(data_source.is_backtesting_mode)  # Default is False for live trading
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_get_historical_prices_success(self, mock_get_data):
-        """Test successful historical price retrieval"""
-        # Create test data
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0, 102.0],
-            'high': [102.0, 103.0, 104.0],
-            'low': [99.0, 100.0, 101.0],
-            'close': [101.0, 102.0, 103.0],
-            'volume': [1000, 1100, 1200]
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00',
-            '2025-01-01 09:32:00'
-        ])
-        
-        mock_get_data.return_value = test_df
-        
-        # Initialize data source
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        # Set current datetime for backtesting
-        data_source._datetime = datetime(2025, 1, 1, 10, 0, 0)
-        
-        # Get historical prices
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=3,
-            timestep="minute"
-        )
-        
-        # Verify result
-        self.assertIsInstance(result, Bars)
-        self.assertEqual(len(result.df), 3)
-        
-        # Verify mock was called with correct parameters
+        self.assertIsInstance(bars, Bars)
+        self.assertEqual(len(bars.df), 3)
         mock_get_data.assert_called_once()
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_get_historical_prices_no_data(self, mock_get_data):
-        """Test historical price retrieval with no data"""
-        mock_get_data.return_value = None
-        
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=10,
-            timestep="minute"
-        )
-        
-        self.assertIsNone(result)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_last_price_from_databento')
-    def test_get_last_price_success(self, mock_get_last_price):
-        """Test successful last price retrieval"""
-        mock_get_last_price.return_value = 4250.75
-        
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        result = data_source.get_last_price(asset=self.test_asset)
-        
-        self.assertEqual(result, 4250.75)
-        mock_get_last_price.assert_called_once_with(
-            api_key=self.api_key,
-            asset=self.test_asset,
-            venue=None
-        )
+    def test_get_historical_prices_returns_none_for_non_futures(self):
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars"
+        ) as mock_get_data:
+            data_source = DataBentoData(**self.datasource_kwargs)
+            result = data_source.get_historical_prices(
+                asset=self.equity_asset,
+                length=5,
+                timestep="minute",
+            )
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_last_price_from_databento')
-    def test_get_last_price_no_data(self, mock_get_last_price):
-        """Test last price retrieval with no data"""
-        mock_get_last_price.return_value = None
-        
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        result = data_source.get_last_price(asset=self.test_asset)
-        
         self.assertIsNone(result)
+        mock_get_data.assert_not_called()
+
+    def test_get_historical_prices_handles_exceptions(self):
+        data_source = DataBentoData(**self.datasource_kwargs)
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            side_effect=RuntimeError("boom"),
+        ):
+            with self.assertRaises(RuntimeError):
+                data_source.get_historical_prices(
+                    asset=self.future_asset,
+                    length=2,
+                    timestep="minute",
+                )
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_get_chains(self):
-        """Test options chains retrieval (should return empty dict)"""
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        result = data_source.get_chains(asset=self.test_asset)
-        
-        self.assertEqual(result, {})
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_get_historical_prices_single_asset(self, mock_get_data):
-        """Test historical price retrieval for a single asset"""
-        # Create test data
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0],
-            'high': [102.0, 103.0],
-            'low': [99.0, 100.0],
-            'close': [101.0, 102.0],
-            'volume': [1000, 1100]
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00'
-        ])
-        
-        mock_get_data.return_value = test_df
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=2,
-            timestep="minute"
-        )
-        
-        self.assertIsNotNone(result)
-        self.assertEqual(len(result.df), 2)
-        self.assertEqual(result.asset, self.test_asset)
-        self.assertEqual(result.source, "DATABENTO")
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_get_historical_prices_error_handling(self, mock_get_data):
-        """Test error handling in get_historical_prices"""
-        # Setup: Mock to raise exception
-        mock_get_data.side_effect = Exception("Test error")
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # This should not raise an exception but return None
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=2,
-            timestep="minute"
-        )
-        
-        self.assertIsNone(result)
+    def test_get_historical_prices_trims_to_requested_length(self):
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            return_value=self._polars_ohlcv(10),
+        ):
+            data_source = DataBentoData(**self.datasource_kwargs)
+            bars = data_source.get_historical_prices(
+                asset=self.future_asset,
+                length=4,
+                timestep="minute",
+            )
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_parse_source_bars_valid_data(self):
-        """Test parsing of valid source data using inherited method"""
-        # Create test DataFrame
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0],
-            'high': [102.0, 103.0],
-            'low': [99.0, 100.0],
-            'close': [101.0, 102.0],
-            'volume': [1000, 1100]
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00'
-        ])
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # Test that _parse_source_bars is available from parent class
-        result = data_source._parse_source_bars({self.test_asset: test_df})
-        
-        self.assertIsInstance(result, dict)
-        self.assertIn(self.test_asset, result)
-        self.assertIsInstance(result[self.test_asset], Bars)
-        self.assertEqual(len(result[self.test_asset].df), 2)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_parse_source_bars_missing_columns(self):
-        """Test parsing of data with missing columns using inherited method"""
-        # Create test DataFrame missing required columns
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0],
-            'high': [102.0, 103.0],
-            # Missing 'low', 'close', 'volume'
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00'
-        ])
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # Test that _parse_source_bars handles missing columns
-        result = data_source._parse_source_bars({self.test_asset: test_df})
-        
-        self.assertIsInstance(result, dict)
-        self.assertIn(self.test_asset, result)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_parse_source_bars_empty_data(self):
-        """Test parsing of empty data using inherited method"""
-        test_df = pd.DataFrame()
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # Test that _parse_source_bars handles empty data
-        result = data_source._parse_source_bars({self.test_asset: test_df})
-        
-        self.assertIsInstance(result, dict)
-        self.assertIn(self.test_asset, result)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_timestep_mapping(self):
-        """Test timestep mapping functionality"""
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # Test valid timestep mappings
-        test_cases = [
-            ("minute", "minute"),
-            ("1m", "minute"),
-            ("hour", "hour"),
-            ("1h", "hour"),
-            ("day", "day"),
-            ("1d", "day"),
-        ]
-        
-        for input_timestep, expected in test_cases:
-            with self.subTest(timestep=input_timestep):
-                result = data_source._parse_source_timestep(input_timestep)
-                self.assertEqual(result, expected)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    def test_backtesting_mode_detection(self):
-        """Test that backtesting mode is properly detected"""
-        # Test with default initialization (live trading mode)
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # DataBento is currently configured for live trading by default
-        self.assertFalse(data_source.is_backtesting_mode)
-        
-        # Test that name and source are set correctly
-        self.assertEqual(data_source.name, "databento")
-        self.assertEqual(data_source.SOURCE, "DATABENTO")
+        self.assertEqual(len(bars.df), 4)
+        self.assertTrue((bars.df.index[-1] > bars.df.index[0]))
+
+    # ------------------------------------------------------------------
+    # Last price & quotes
+    # ------------------------------------------------------------------
+    def test_get_last_price_uses_historical_fallback(self):
+        frame = self._polars_ohlcv(2)
+        last_close = float(frame.select("close").to_series().tail(1)[0])
+
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            return_value=frame,
+        ):
+            data_source = DataBentoData(**self.datasource_kwargs)
+            price = data_source.get_last_price(asset=self.future_asset)
+
+        self.assertEqual(price, last_close)
+
+    def test_get_last_price_returns_none_when_no_data(self):
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            return_value=None,
+        ):
+            data_source = DataBentoData(**self.datasource_kwargs)
+            price = data_source.get_last_price(asset=self.future_asset)
+
+        self.assertIsNone(price)
+
+    def test_get_quote_falls_back_to_last_price(self):
+        with patch.object(DataBentoData, "get_last_price", return_value=123.45):
+            data_source = DataBentoData(**self.datasource_kwargs)
+            quote = data_source.get_quote(asset=self.future_asset)
+
+        self.assertEqual(quote.asset, self.future_asset)
+        self.assertEqual(quote.price, 123.45)
+        self.assertGreaterEqual(quote.ask, quote.bid)
+
+    # ------------------------------------------------------------------
+    # Continuous futures resolution
+    # ------------------------------------------------------------------
+    def test_continuous_future_resolves_symbol(self):
+        with patch(
+            "lumibot.data_sources.databento_data_polars.databento_helper_polars.get_price_data_from_databento_polars",
+            return_value=self._polars_ohlcv(2),
+        ):
+            data_source = DataBentoData(**self.datasource_kwargs)
+            bars = data_source.get_historical_prices(
+                asset=self.cont_future_asset,
+                length=2,
+                timestep="minute",
+            )
 
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_get_historical_prices_backtesting_path(self, mock_get_data):
-        """Test that get_historical_prices uses backtesting path when in backtesting mode"""
-        # Create test data
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0, 102.0],
-            'high': [102.0, 103.0, 104.0],
-            'low': [99.0, 100.0, 101.0],
-            'close': [101.0, 102.0, 103.0],
-            'volume': [1000, 1100, 1200]
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00',
-            '2025-01-01 09:32:00'
-        ])
-        
-        mock_get_data.return_value = test_df
-        
-        data_source = DataBentoData(
-            api_key=self.api_key,
-            datetime_start=self.start_date,
-            datetime_end=self.end_date
-        )
-        
-        # Test that the method works correctly
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=3,
-            timestep="minute"
-        )
-        
-        # Verify that the data was retrieved
-        self.assertIsNotNone(result)
-        self.assertIsInstance(result, Bars)
-        self.assertEqual(len(result.df), 3)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_timezone_handling(self, mock_get_data):
-        """Test timezone handling in get_historical_prices"""
-        # Create test data
-        test_df = pd.DataFrame({
-            'open': [100.0, 101.0],
-            'high': [102.0, 103.0],
-            'low': [99.0, 100.0],
-            'close': [101.0, 102.0],
-            'volume': [1000, 1100]
-        })
-        test_df.index = pd.to_datetime([
-            '2025-01-01 09:30:00',
-            '2025-01-01 09:31:00'
-        ])
-        
-        mock_get_data.return_value = test_df
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # This should not raise an exception
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=2,
-            timestep="minute"
-        )
-        
-        self.assertIsNotNone(result)
-        self.assertIsInstance(result, Bars)
-        self.assertEqual(len(result.df), 2)
-
-    @patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True)
-    @patch('lumibot.tools.databento_helper.get_price_data_from_databento')
-    def test_error_handling(self, mock_get_data):
-        """Test error handling in get_historical_prices"""
-        # Setup: Mock to raise exception
-        mock_get_data.side_effect = Exception("Test error")
-        
-        data_source = DataBentoData(api_key=self.api_key)
-        
-        # This should not raise an exception but return None
-        result = data_source.get_historical_prices(
-            asset=self.test_asset,
-            length=1,
-            timestep="minute"
-        )
-        
-        # Should return None on error
-        self.assertIsNone(result)
+        self.assertIsNotNone(bars)
+        self.assertEqual(bars.asset, self.cont_future_asset)
 
+    # ------------------------------------------------------------------
+    # Integration-style helpers (mocked)
+    # ------------------------------------------------------------------
     def test_environment_dates_integration(self):
-        """Test DataBento with environment file dates"""
         from dotenv import load_dotenv
-        import os
-        
-        # Load environment variables from the strategy .env file
+
         env_path = "/Users/robertgrzesik/Documents/Development/Strategy Library/Alligator Futures Bot Strategy/src/.env"
         if os.path.exists(env_path):
             load_dotenv(env_path)
-        
-        # Get dates from environment variables
-        start_str = os.getenv("BACKTESTING_START", "2024-01-01")
-        end_str = os.getenv("BACKTESTING_END", "2024-12-31")
-        
-        start_date = datetime.strptime(start_str, "%Y-%m-%d")
-        end_date = datetime.strptime(end_str, "%Y-%m-%d")
-        
-        # Use a recent subset for testing (last 3 months)
-        test_end_date = datetime(2024, 12, 31)
-        test_start_date = datetime(2024, 10, 1)
-        
-        with patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True):
-            data_source = DataBentoData(api_key=self.api_key)
-            
-            # Test with MES continuous futures
-            mes_asset = Asset("MES", asset_type=Asset.AssetType.CONT_FUTURE)
-            
-            # Mock the get_historical_prices method
-            mock_bars = Mock()
-            mock_bars.df = pd.DataFrame({
-                'open': [4500, 4505, 4510],
-                'high': [4510, 4515, 4520],
-                'low': [4495, 4500, 4505],
-                'close': [4505, 4510, 4515],
-                'volume': [1000, 1100, 1200]
-            }, index=pd.date_range(start=test_start_date, periods=3, freq='h'))
-            
-            with patch.object(data_source, 'get_historical_prices', return_value=mock_bars):
-                bars = data_source.get_historical_prices(
-                    asset=mes_asset,
-                    length=60,
-                    timestep="minute"
-                )
-                
-                self.assertIsNotNone(bars)
-                self.assertIsNotNone(bars.df)
-                self.assertEqual(len(bars.df), 3)
-    
+
+        with patch.object(DataBentoData, "get_historical_prices", return_value=self._bars(3)) as mock_get_hist:
+            data_source = DataBentoData(**self.datasource_kwargs)
+            bars = data_source.get_historical_prices(
+                asset=self.cont_future_asset,
+                length=60,
+                timestep="minute",
+            )
+
+        mock_get_hist.assert_called_once()
+        self.assertIsNotNone(bars)
+        self.assertEqual(len(bars.df), 3)
+
     def test_mes_strategy_logic_simulation(self):
-        """Test MES strategy logic with simulated data"""
-        with patch('lumibot.tools.databento_helper.DATABENTO_AVAILABLE', True):
-            data_source = DataBentoData(api_key=self.api_key)
-            
-            mes_asset = Asset("MES", asset_type=Asset.AssetType.CONT_FUTURE)
-            
-            # Create mock data that simulates 60 minutes of MES futures
-            mock_bars = Mock()
-            mock_bars.df = pd.DataFrame({
-                'open': [4500 + i for i in range(60)],
-                'high': [4510 + i for i in range(60)],
-                'low': [4490 + i for i in range(60)],
-                'close': [4505 + i for i in range(60)],
-                'volume': [1000 + i*10 for i in range(60)]
-            }, index=pd.date_range(start=datetime(2024, 6, 10, 8, 0), periods=60, freq='min'))
-            
-            with patch.object(data_source, 'get_historical_prices', return_value=mock_bars):
-                bars = data_source.get_historical_prices(
-                    asset=mes_asset,
-                    length=60,
-                    timestep="minute"
-                )
-                
-                self.assertIsNotNone(bars)
-                self.assertIsNotNone(bars.df)
-                self.assertEqual(len(bars.df), 60)
-                
-                # Test strategy logic
-                df = bars.df
-                current_price = df["close"].iloc[-1]
-                sma_60 = df["close"].mean()
-                
-                # Should have a clear trend in our test data
-                self.assertGreater(current_price, sma_60)
-                self.assertGreater(current_price, 4500)  # Should be trending up
-
-if __name__ == '__main__':
+        data_source = DataBentoData(**self.datasource_kwargs)
+        mock_bars = MagicMock()
+        mock_df = pd.DataFrame(
+            {
+                "open": [4500 + i for i in range(60)],
+                "high": [4510 + i for i in range(60)],
+                "low": [4490 + i for i in range(60)],
+                "close": [4505 + i for i in range(60)],
+                "volume": [1_000 + i * 10 for i in range(60)],
+            },
+            index=pd.date_range(start=datetime(2024, 6, 10, 8, 0), periods=60, freq="min"),
+        )
+        mock_bars.df = mock_df
+
+        with patch.object(data_source, "get_historical_prices", return_value=mock_bars):
+            bars = data_source.get_historical_prices(
+                asset=self.cont_future_asset,
+                length=60,
+                timestep="minute",
+            )
+
+        self.assertEqual(len(bars.df), 60)
+        current_price = bars.df["close"].iloc[-1]
+        sma_60 = bars.df["close"].mean()
+
+        self.assertGreater(current_price, sma_60)
+        self.assertGreater(current_price, 4500)
+
+
+if __name__ == "__main__":
     unittest.main()