121 lines
5.0 KiB
Python
121 lines
5.0 KiB
Python
import pytest
|
|
import pandas as pd
|
|
from unittest.mock import Mock, patch
|
|
from datetime import datetime
|
|
|
|
from components.charts.data_integration import MarketDataIntegrator
|
|
from components.charts.indicator_manager import IndicatorManager
|
|
from components.charts.layers.indicators import IndicatorLayerConfig
|
|
|
|
@pytest.fixture
|
|
def market_data_integrator_components():
|
|
"""Provides a complete setup for testing MarketDataIntegrator."""
|
|
|
|
# 1. Main DataFrame (e.g., 1h)
|
|
main_timestamps = pd.to_datetime(['2024-01-01 10:00', '2024-01-01 11:00', '2024-01-01 12:00', '2024-01-01 13:00'], utc=True)
|
|
main_df = pd.DataFrame({'close': [100, 102, 101, 103]}, index=main_timestamps)
|
|
|
|
# 2. Higher-timeframe DataFrame (e.g., 4h)
|
|
indicator_timestamps = pd.to_datetime(['2024-01-01 08:00', '2024-01-01 12:00'], utc=True)
|
|
indicator_df_raw = [{'timestamp': ts, 'close': val} for ts, val in zip(indicator_timestamps, [98, 101.5])]
|
|
|
|
# 3. Mock IndicatorManager and configs
|
|
indicator_manager = Mock(spec=IndicatorManager)
|
|
user_indicator = Mock()
|
|
user_indicator.id = 'rsi_4h'
|
|
user_indicator.name = 'RSI'
|
|
user_indicator.timeframe = '4h'
|
|
user_indicator.type = 'rsi'
|
|
user_indicator.parameters = {'period': 14}
|
|
|
|
indicator_manager.load_indicator.return_value = user_indicator
|
|
|
|
indicator_config = Mock(spec=IndicatorLayerConfig)
|
|
indicator_config.id = 'rsi_4h'
|
|
|
|
# 4. DataIntegrator instance
|
|
integrator = MarketDataIntegrator()
|
|
|
|
# 5. Mock internal fetching and calculation
|
|
# Mock get_market_data_for_indicators to return raw candles
|
|
integrator.get_market_data_for_indicators = Mock(return_value=(indicator_df_raw, []))
|
|
|
|
# Mock indicator calculation result
|
|
indicator_result_values = [{'timestamp': indicator_timestamps[1], 'rsi': 55.0}] # Only one valid point
|
|
indicator_pkg = {'data': [Mock(timestamp=r['timestamp'], values={'rsi': r['rsi']}) for r in indicator_result_values]}
|
|
integrator.indicators.calculate = Mock(return_value=indicator_pkg)
|
|
|
|
return integrator, main_df, indicator_config, indicator_manager, user_indicator
|
|
|
|
def test_multi_timeframe_alignment(market_data_integrator_components):
|
|
"""
|
|
Tests that indicator data from a higher timeframe is correctly aligned
|
|
with the main chart's data.
|
|
"""
|
|
integrator, main_df, indicator_config, indicator_manager, user_indicator = market_data_integrator_components
|
|
|
|
# Execute the method to test
|
|
indicator_data_map = integrator.get_indicator_data(
|
|
main_df=main_df,
|
|
main_timeframe='1h',
|
|
indicator_configs=[indicator_config],
|
|
indicator_manager=indicator_manager,
|
|
symbol='BTC-USDT'
|
|
)
|
|
|
|
# --- Assertions ---
|
|
assert user_indicator.id in indicator_data_map
|
|
aligned_data = indicator_data_map[user_indicator.id]
|
|
|
|
# Expected series after reindexing and forward-filling
|
|
expected_series = pd.Series(
|
|
[None, None, 55.0, 55.0],
|
|
index=main_df.index,
|
|
name='rsi'
|
|
)
|
|
|
|
result_series = aligned_data['rsi']
|
|
pd.testing.assert_series_equal(result_series, expected_series, check_index_type=False)
|
|
|
|
@patch('components.charts.utils.prepare_chart_data', lambda x: pd.DataFrame(x).set_index('timestamp'))
|
|
def test_no_custom_timeframe_uses_main_df(market_data_integrator_components):
|
|
"""
|
|
Tests that if an indicator has no custom timeframe, it uses the main
|
|
DataFrame for calculation.
|
|
"""
|
|
integrator, main_df, indicator_config, indicator_manager, user_indicator = market_data_integrator_components
|
|
|
|
# Override indicator to have no timeframe
|
|
user_indicator.timeframe = None
|
|
indicator_manager.load_indicator.return_value = user_indicator
|
|
|
|
# Mock calculation result on main_df
|
|
result_timestamps = main_df.index[1:]
|
|
indicator_result_values = [{'timestamp': ts, 'sma': val} for ts, val in zip(result_timestamps, [101.0, 101.5, 102.0])]
|
|
indicator_pkg = {'data': [Mock(timestamp=r['timestamp'], values={'sma': r['sma']}) for r in indicator_result_values]}
|
|
integrator.indicators.calculate = Mock(return_value=indicator_pkg)
|
|
|
|
# Execute
|
|
indicator_data_map = integrator.get_indicator_data(
|
|
main_df=main_df,
|
|
main_timeframe='1h',
|
|
indicator_configs=[indicator_config],
|
|
indicator_manager=indicator_manager,
|
|
symbol='BTC-USDT'
|
|
)
|
|
|
|
# Assert that get_market_data_for_indicators was NOT called
|
|
integrator.get_market_data_for_indicators.assert_not_called()
|
|
|
|
# Assert that calculate was called with main_df
|
|
integrator.indicators.calculate.assert_called_with('rsi', main_df, period=14)
|
|
|
|
# Assert the result is what we expect
|
|
assert user_indicator.id in indicator_data_map
|
|
result_series = indicator_data_map[user_indicator.id]['sma']
|
|
expected_series = pd.Series([101.0, 101.5, 102.0], index=result_timestamps, name='sma')
|
|
|
|
# Reindex expected to match the result's index for comparison
|
|
expected_series = expected_series.reindex(main_df.index)
|
|
|
|
pd.testing.assert_series_equal(result_series, expected_series, check_index_type=False) |