diff --git a/tests/pipeline/test_technical.py b/tests/pipeline/test_technical.py index c2b1dfc6..313337e4 100644 --- a/tests/pipeline/test_technical.py +++ b/tests/pipeline/test_technical.py @@ -72,17 +72,17 @@ class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase): middle_cols = [] upper_cols = [] for n in range(self.nassets): - try: + close_col = closes[:, n] + if np.isnan(close_col).all(): + # ta-lib doesn't deal well with all nans. + upper, middle, lower = [np.full(self.ndays, np.nan)] * 3 + else: upper, middle, lower = talib.BBANDS( - closes[:, n], + close_col, window_length, k, k, ) - except Exception: - # If the input array is all nan then talib raises an instance - # of Exception. - upper, middle, lower = [np.full(self.ndays, np.nan)] * 3 upper_cols.append(upper) middle_cols.append(middle) @@ -128,15 +128,13 @@ class BollingerBandsTestCase(WithTechnicalFactor, ZiplineTestCase): closes, ) - assert_equal( - result.upper, - expected_upper, - ) - assert_equal( - result.middle, - expected_middle, - ) - assert_equal( - result.lower, - expected_lower, - ) + assert_equal(result.upper, expected_upper) + assert_equal(result.middle, expected_middle) + assert_equal(result.lower, expected_lower) + + def test_bollinger_bands_output_ordering(self): + bbands = BollingerBands(window_length=5, k=2) + lower, middle, upper = bbands + self.assertIs(lower, bbands.lower) + self.assertIs(middle, bbands.middle) + self.assertIs(upper, bbands.upper)