cfdaily
196 lines
5.8 KiB
Python
196 lines
5.8 KiB
Python
"""
|
||
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[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)
|
||
|
||
|
||
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'])
|