new api in trading_dates_mixin and test_case

ricequant · Nov 12, 2018 · 0a34912 · 0a34912
1 parent a7688dc
commit 0a34912
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 2 deletions.
diff --git a/rqalpha/data/trading_dates_mixin.py b/rqalpha/data/trading_dates_mixin.py
@@ -85,3 +85,10 @@ def get_n_trading_dates_until(self, dt, n):
             return self._dates[pos - n:pos]
 
         return self._dates[:pos]
+
+    def count_trading_dates(self, start_date, end_date):
+        start_date = _to_timestamp(start_date)
+        end_date = _to_timestamp(end_date)
+
+        return self._dates.searchsorted(end_date, side='right') - self._dates.searchsorted(start_date)
+
diff --git a/rqalpha/utils/testing/fixtures.py b/rqalpha/utils/testing/fixtures.py
@@ -63,7 +63,6 @@ def init_fixture(self):
         else:
             from backports.tempfile import TemporaryDirectory
 
-
         super(TempDirFixture, self).init_fixture()
         self.temp_dir = TemporaryDirectory()
 
@@ -122,6 +121,7 @@ def init_fixture(self):
 class DataProxyFixture(BaseDataSourceFixture):
     def __init__(self, *args, **kwargs):
         super(DataProxyFixture, self).__init__(*args, **kwargs)
+        self.data_proxy = None
         self.data_source = None
 
     def init_fixture(self):
@@ -130,7 +130,8 @@ def init_fixture(self):
         super(DataProxyFixture, self).init_fixture()
         if not self.data_source:
             self.data_source = self.base_data_source
-        self.env.set_data_proxy(DataProxy(self.data_source))
+        self.data_proxy = DataProxy(self.data_source)
+        self.env.set_data_proxy(self.data_proxy)
 
     @contextmanager
     def mock_data_proxy_method(self, name, mock_method):

diff --git a/tests/unittest/test_data/test_trading_dates_mixin.py b/tests/unittest/test_data/test_trading_dates_mixin.py
@@ -0,0 +1,13 @@
+from rqalpha.utils.testing import DataProxyFixture, RQAlphaTestCase
+
+
+class TradingDateMixinTestCase(DataProxyFixture, RQAlphaTestCase):
+    def init_fixture(self):
+        super(TradingDateMixinTestCase, self).init_fixture()
+
+    def test_count_trading_dates(self):
+        from datetime import date
+
+        assert self.data_proxy.count_trading_dates(date(2018, 11, 1), date(2018, 11, 12)) == 8
+        assert self.data_proxy.count_trading_dates(date(2018, 11, 3), date(2018, 11, 12)) == 6
+        assert self.data_proxy.count_trading_dates(date(2018, 11, 3), date(2018, 11, 18)) == 10