diff --git a/zhangfei-technical/research/task-20260417-atr-indicator/test_atr.py b/zhangfei-technical/research/task-20260417-atr-indicator/test_atr.py
index 087d1f6a6..ebad95db0 100644
--- a/zhangfei-technical/research/task-20260417-atr-indicator/test_atr.py
+++ b/zhangfei-technical/research/task-20260417-atr-indicator/test_atr.py
@@ -167,7 +167,8 @@ def test_calculate_atr_convenience():
     
     assert 'TR' in result.columns
     assert 'ATR_5' in result.columns
-    assert result.shape == df.shape
+    assert result.shape[0] == df.shape[0]
+    assert result.shape[1] == df.shape[1] + 2
     
     result_dropna = calculate_atr(df, period=5, method='sma', drop_na=True)
     assert len(result_dropna) < len(df)