From ea96abecd06b753f0780b1875dc92c4b884308db Mon Sep 17 00:00:00 2001
From: Joe Schr <8218910+TheJoeSchr@users.noreply.github.com>
Date: Fri, 3 Jan 2025 16:27:18 +0100
Subject: [PATCH] test: Add comprehensive test for stacked_imbalances with
 multiple price entries

---
 tests/data/test_converter_orderflow.py | 36 ++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)

diff --git a/tests/data/test_converter_orderflow.py b/tests/data/test_converter_orderflow.py
index acef1ef4f..50b47b7d5 100644
--- a/tests/data/test_converter_orderflow.py
+++ b/tests/data/test_converter_orderflow.py
@@ -7,6 +7,7 @@ from freqtrade.data.converter.orderflow import (
     ORDERFLOW_ADDED_COLUMNS,
     timeframe_to_DateOffset,
     trades_to_volumeprofile_with_total_delta_bid_ask,
+    stacked_imbalance,
 )
 from freqtrade.data.converter.trade_converter import trades_list_to_df
 from freqtrade.data.dataprovider import DataProvider
@@ -569,6 +570,41 @@ def test_analyze_with_orderflow(
     assert isinstance(lastval_of2, dict)
 
 
+def test_stacked_imbalances_multiple_prices():
+    """Test that stacked imbalances correctly returns multiple price levels when present"""
+    # Test with empty result
+    df_no_stacks = pd.DataFrame(
+        {
+            'bid_imbalance': [False, False, True, False],
+            'ask_imbalance': [False, True, False, False]
+        },
+        index=[234.95, 234.96, 234.97, 234.98]
+    )
+    no_stacks = stacked_imbalance(df_no_stacks, "bid", stacked_imbalance_range=2, should_reverse=False)
+    assert no_stacks == [np.nan]
+    
+    # Create a sample DataFrame with known imbalances
+    df = pd.DataFrame(
+        {
+            'bid_imbalance': [True, True, True, False, False, True, True, False],
+            'ask_imbalance': [False, False, True, True, True, False, False, True]
+        },
+        index=[234.95, 234.96, 234.97, 234.98, 234.99, 235.00, 235.01, 235.02]
+    )
+    # Test bid imbalances (should return prices in ascending order)
+    bid_prices = stacked_imbalance(df, "bid", stacked_imbalance_range=2, should_reverse=False)
+    assert bid_prices == [234.95, 234.96, 234.97, 235.00, 235.01]
+    
+    # Test ask imbalances (should return prices in descending order)
+    ask_prices = stacked_imbalance(df, "ask", stacked_imbalance_range=2, should_reverse=True)
+    assert ask_prices == [235.02, 234.99, 234.98, 234.97]
+    
+    # Test with higher stacked_imbalance_range
+    bid_prices_higher = stacked_imbalance(df, "bid", stacked_imbalance_range=3, should_reverse=False)
+    assert bid_prices_higher == [234.95, 234.96, 234.97]
+    
+
+
 def test_timeframe_to_DateOffset():
     assert timeframe_to_DateOffset("1s") == pd.DateOffset(seconds=1)
     assert timeframe_to_DateOffset("1m") == pd.DateOffset(minutes=1)