From 4963b2ca72cc0133a050f5960152ff12272d2efd Mon Sep 17 00:00:00 2001
From: llllllllll <joe@quantopian.com>
Date: Thu, 10 Dec 2015 15:14:15 -0500
Subject: [PATCH] TST: Adds tests for infer_timestamp

---
 tests/pipeline/test_earnings.py      | 30 ++++++++++++++++++++++++----
 zipline/pipeline/loaders/earnings.py | 14 ++++++-------
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/tests/pipeline/test_earnings.py b/tests/pipeline/test_earnings.py
index 32509005..871fb155 100644
--- a/tests/pipeline/test_earnings.py
+++ b/tests/pipeline/test_earnings.py
@@ -4,6 +4,7 @@ Tests for the reference loader for EarningsCalendar.
 from unittest import TestCase
 
 import blaze as bz
+from blaze.compute.core import swap_resources_into_scope
 from contextlib2 import ExitStack
 from nose_parameterized import parameterized
 import pandas as pd
@@ -388,9 +389,30 @@ class BlazeEarningsCalendarLoaderNotInteractiveTestCase(
             BlazeEarningsCalendarLoaderNotInteractiveTestCase,
             self,
         ).loader_args(dates)
+        return swap_resources_into_scope(bound_expr, {})
 
-        expr = bz.symbol('expr', bound_expr.dshape)
-        return (
-            expr,
-            {expr: bound_expr._resources()[bound_expr]},
+
+class EarningsCalendarLoaderInferTimestampTestCase(TestCase):
+    def test_infer_timestamp(self):
+        dtx = pd.date_range('2014-01-01', '2014-01-10')
+        announcement_dates = {
+            0: dtx,
+            1: pd.Series(dtx, dtx),
+        }
+        loader = EarningsCalendarLoader(
+            dtx,
+            announcement_dates,
+            infer_timestamps=True,
+        )
+        self.assertEqual(
+            loader.announcement_dates.keys(),
+            announcement_dates.keys(),
+        )
+        assert_series_equal(
+            loader.announcement_dates[0],
+            pd.Series(index=[dtx[0]] * 10, data=dtx),
+        )
+        assert_series_equal(
+            loader.announcement_dates[1],
+            announcement_dates[1],
         )
diff --git a/zipline/pipeline/loaders/earnings.py b/zipline/pipeline/loaders/earnings.py
index acde30ef..cb8106b6 100644
--- a/zipline/pipeline/loaders/earnings.py
+++ b/zipline/pipeline/loaders/earnings.py
@@ -45,12 +45,12 @@ class EarningsCalendarLoader(PipelineLoader):
         ``announcement_dates``.
     """
     def __init__(self, all_dates, announcement_dates, infer_timestamps=False):
-        self._all_dates = all_dates
+        self.all_dates = all_dates
 
-        self._announcement_dates = announcement_dates = (
+        self.announcement_dates = announcement_dates = (
             announcement_dates.copy()
         )
-        dates = self._all_dates.values
+        dates = self.all_dates.values
         for k, v in iteritems(announcement_dates):
             if isinstance(v, pd.DatetimeIndex):
                 if not infer_timestamps:
@@ -80,8 +80,8 @@ class EarningsCalendarLoader(PipelineLoader):
         return DataFrameLoader(
             EarningsCalendar.next_announcement,
             next_earnings_date_frame(
-                self._all_dates,
-                self._announcement_dates,
+                self.all_dates,
+                self.announcement_dates,
             ),
             adjustments=None,
         )
@@ -91,8 +91,8 @@ class EarningsCalendarLoader(PipelineLoader):
         return DataFrameLoader(
             EarningsCalendar.previous_announcement,
             previous_earnings_date_frame(
-                self._all_dates,
-                self._announcement_dates,
+                self.all_dates,
+                self.announcement_dates,
             ),
             adjustments=None,
         )