def test_select_n(): algo = algos.SelectN(n=1, sort_descending=True) s = bt.Strategy('s') dts = pd.date_range('2010-01-01', periods=3) data = pd.DataFrame(index=dts, columns=['c1', 'c2'], data=100.) data['c1'].ix[dts[2]] = 105 data['c2'].ix[dts[2]] = 95 s.setup(data) s.update(dts[2]) s.temp['stat'] = data.calc_total_return() assert algo(s) selected = s.temp['selected'] assert len(selected) == 1 assert 'c1' in selected algo = algos.SelectN(n=1, sort_descending=False) assert algo(s) selected = s.temp['selected'] assert len(selected) == 1 assert 'c2' in selected # return 2 we have if all_or_none false algo = algos.SelectN(n=3, sort_descending=False) assert algo(s) selected = s.temp['selected'] assert len(selected) == 2 assert 'c1' in selected assert 'c2' in selected # return 0 we have if all_or_none true algo = algos.SelectN(n=3, sort_descending=False, all_or_none=True) assert algo(s) selected = s.temp['selected'] assert len(selected) == 0
def test_select_n_perc(): algo = algos.SelectN(n=0.5, sort_descending=True) s = bt.Strategy('s') dts = pd.date_range('2010-01-01', periods=3) data = pd.DataFrame(index=dts, columns=['c1', 'c2'], data=100.) data['c1'].ix[dts[2]] = 105 data['c2'].ix[dts[2]] = 95 s.setup(data) s.update(dts[2]) s.temp['stat'] = data.calc_total_return() assert algo(s) selected = s.temp['selected'] assert len(selected) == 1 assert 'c1' in selected