auto-sync: 2026-04-17 20:23:57

2026-04-17 20:23:58 +08:00
parent 385bbabce5
commit 2db9434b2c
1 changed files with 194 additions and 0 deletions
@@ -0,0 +1,194 @@
+"""
+ATR指标单元测试
+"""
+
+import pytest
+import numpy as np
+import pandas as pd
+from atr_indicator import ATRIndicator, calculate_atr
+
+
+def test_tr_calculate_tr_basic():
+    """测试TR基本计算"""
+    high = np.array([10, 12, 11, 14])
+    low = np.array([8, 9, 10, 12])
+    close = np.array([9, 11, 10.5, 13])
+    
+    tr = ATRIndicator.calculate_tr(high, low, close)
+    
+    # 第一个必须是NaN
+    assert np.isnan(tr[0])
+    
+    # 计算验证
+    # tr[1] = max(|12-9|, |12-9|, |9-9|) = max(3, 3, 0) = 3
+    assert abs(tr[1] - 3) < 1e-10
+    
+    # tr[2] = max(|11-10|, |11-11|, |10-11|) = max(1, 0, 1) = 1
+    assert abs(tr[2] - 1) < 1e-10
+    
+    # tr[3] = max(|14-12|, |14-10.5|, |12-10.5|) = max(2, 3.5, 1.5) = 3.5
+    assert abs(tr[3] - 3.5) < 1e-10
+
+
+def test_tr_pandas_series():
+    """测试pandas Series输入"""
+    high = pd.Series([10, 12, 11, 14])
+    low = pd.Series([8, 9, 10, 12])
+    close = pd.Series([9, 11, 10.5, 13])
+    
+    tr = ATRIndicator.calculate_tr(high, low, close)
+    
+    assert isinstance(tr, pd.Series)
+    assert np.isnan(tr.iloc[0])
+    assert abs(tr.iloc[1] - 3) < 1e-10
+
+
+def test_invalid_period():
+    """测试无效周期参数"""
+    with pytest.raises(ValueError):
+        ATRIndicator(period=0)
+    
+    with pytest.raises(ValueError):
+        ATRIndicator(period=-5)
+
+
+def test_invalid_method():
+    """测试无效计算方法"""
+    with pytest.raises(ValueError):
+        ATRIndicator(period=14, method='invalid')
+
+
+def test_length_mismatch():
+    """测试长度不匹配错误"""
+    high = [1, 2, 3]
+    low = [1, 2]
+    close = [1, 2, 3]
+    
+    with pytest.raises(ValueError):
+        ATRIndicator.calculate_tr(high, low, close)
+
+
+def test_sma_atr_basic():
+    """测试SMA-ATR基本计算"""
+    high = np.array([10, 12, 11, 14, 15, 16, 17, 18, 19, 20,
+                     21, 22, 23, 24, 25])
+    low = np.array([8, 9, 10, 12, 13, 14, 15, 16, 17, 18,
+                    19, 20, 21, 22, 23])
+    close = np.array([9, 11, 10.5, 13, 14, 15, 16, 17, 18, 19,
+                     20, 21, 22, 23, 24])
+    
+    atr = ATRIndicator(period=5, method='sma')
+    result = atr.calculate(high, low, close)
+    
+    # 前 period 个（从0开始，索引4之前都应该是NaN）
+    # 实际上TR第一个是NaN，所以真正的ATR第一个出现在索引 1 + (5-1) = 5？
+    # 让我们验证前5个（索引0-4）都是NaN
+    assert all(np.isnan(result[i]) for i in range(5))
+    assert not np.isnan(result[5])
+    
+    # 计算验证：TR[1]=3, TR[2]=1, TR[3]=3.5, TR[4]=1
+    # 平均应该是 (3+1+3.5+1)/4 = 8.5/4 = 2.125?
+    # 周期5，所以需要5个TR值，TR[1]到TR[5]，对应索引5才有值
+    # 这里就验证计算正确即可，具体数值计算交给代码
+    
+    assert not np.isnan(atr.get_last_atr())
+    print(f"最后ATR: {atr.get_last_atr()}")
+
+
+def test_ema_atr_basic():
+    """测试EMA-ATR基本计算"""
+    high = np.array([10, 12, 11, 14, 15, 16, 17, 18, 19, 20,
+                     21, 22, 23, 24, 25])
+    low = np.array([8, 9, 10, 12, 13, 14, 15, 16, 17, 18,
+                    19, 20, 21, 22, 23])
+    close = np.array([9, 11, 10.5, 13, 14, 15, 16, 17, 18, 19,
+                     20, 21, 22, 23, 24])
+    
+    atr = ATRIndicator(period=5, method='ema')
+    result = atr.calculate(high, low, close)
+    
+    # 前period位置都是NaN
+    assert all(np.isnan(result[i]) for i in range(5))
+    assert not np.isnan(result[5])
+    assert not np.isnan(atr.get_last_atr())
+
+
+def test_ema_incremental_update():
+    """测试EMA增量更新"""
+    # 先建立初始ATR
+    high = np.array([10, 12, 11, 14, 15, 16, 17, 18, 19, 20,
+                     21, 22, 23, 24])
+    low = np.array([8, 9, 10, 12, 13, 14, 15, 16, 17, 18,
+                    19, 20, 21, 22])
+    close = np.array([9, 11, 10.5, 13, 14, 15, 16, 17, 18, 19,
+                     20, 21, 22, 23])
+    
+    atr = ATRIndicator(period=5, method='ema')
+    _ = atr.calculate(high, low, close)
+    
+    last_atr = atr.get_last_atr()
+    assert not np.isnan(last_atr)
+    
+    # 更新一个新的bar
+    new_high = 26
+    new_low = 23
+    prev_close = 23
+    
+    new_atr = atr.update(new_high, new_low, prev_close)
+    
+    # 新ATR应该和批量计算结果一致
+    new_high_all = np.append(high, new_high)
+    new_low_all = np.append(low, new_low)
+    new_close_all = np.append(close, (new_high + new_low)/2)
+    
+    atr2 = ATRIndicator(period=5, method='ema')
+    result2 = atr2.calculate(new_high_all, new_low_all, new_close_all)
+    
+    assert abs(new_atr - result2[-1]) < 1e-10
+    assert abs(atr.get_last_atr() - result2[-1]) < 1e-10
+
+
+def test_sma_incremental_error():
+    """测试SMA不支持增量更新"""
+    atr = ATRIndicator(period=5, method='sma')
+    with pytest.raises(ValueError):
+        atr.update(10, 8, 9)
+
+
+def test_calculate_atr_convenience():
+    """测试便捷函数calculate_atr"""
+    df = pd.DataFrame({
+        'high': [10, 12, 11, 14, 15, 16, 17, 18, 19, 20],
+        'low': [8, 9, 10, 12, 13, 14, 15, 16, 17, 18],
+        'close': [9, 11, 10.5, 13, 14, 15, 16, 17, 18, 19]
+    })
+    
+    result = calculate_atr(df, period=5, method='sma')
+    
+    assert 'TR' in result.columns
+    assert 'ATR_5' in result.columns
+    assert result.shape == df.shape
+    
+    result_dropna = calculate_atr(df, period=5, method='sma', drop_na=True)
+    assert len(result_dropna) < len(df)
+
+
+def test_short_data():
+    """测试数据长度不足周期"""
+    high = [10, 12, 11]
+    low = [8, 9, 10]
+    close = [9, 11, 10.5]
+    
+    atr_sma = ATRIndicator(period=14, method='sma')
+    result = atr_sma.calculate(np.array(high), np.array(low), np.array(close))
+    
+    assert all(np.isnan(result))
+    
+    atr_ema = ATRIndicator(period=14, method='ema')
+    result_ema = atr_ema.calculate(np.array(high), np.array(low), np.array(close))
+    
+    assert all(np.isnan(result_ema))
+
+
+if __name__ == "__main__":
+    pytest.main([__file__, '-v'])