Add unit tests around the ranges that filter will encounter

Author: bmoscon

arctic/chunkstore/chunkstore.py

@@ -85,7 +85,7 @@ def delete(self, symbol, chunk_range=None):
         if chunk_range:
             # read out chunks that fall within the range and filter out
             # data within the range
-            df = self.read(symbol, chunk_range=chunk_range, no_filter=True)
+            df = self.read(symbol, chunk_range=chunk_range, filter_data=False)
             df = self.chunker.exclude(df, chunk_range)
 
             # remove chunks, and update any remaining data
@@ -112,7 +112,7 @@ def list_symbols(self):
     def _get_symbol_info(self, symbol):
         return self._symbols.find_one({'symbol': symbol})
 
-    def read(self, symbol, chunk_range=None, no_filter=False):
+    def read(self, symbol, chunk_range=None, filter_data=True):
         """
         Reads data for a given symbol from the database.
 
@@ -123,7 +123,7 @@ def read(self, symbol, chunk_range=None, no_filter=False):
         chunk_range: object
             corresponding range object for the specified chunker (for
             DateChunker it is a DateRange object)
-        no_filter: boolean
+        filter: boolean
             perform chunk level filtering on the data (see filter() in _chunker)
             only applicable when chunk_range is specified
 
@@ -153,7 +153,7 @@ def read(self, symbol, chunk_range=None, no_filter=False):
 
         data = deserialize(records, sym['type'])
 
-        if no_filter or chunk_range is None:
+        if not filter_data or chunk_range is None:
             return data
         return self.chunker.filter(data, chunk_range)
 

tests/unit/chunkstore/__init__.py

@@ -0,0 +1,30 @@
+from arctic.chunkstore.date_chunker import DateChunker
+from pandas import DataFrame, MultiIndex
+from datetime import datetime as dt
+from arctic.date import DateRange
+from pandas.util.testing import assert_frame_equal
+
+
+def test_date_filter():
+    c = DateChunker()
+    df = DataFrame(data={'data': [1, 2, 3]},
+                   index=MultiIndex.from_tuples([(dt(2016, 1, 1), 1),
+                                                 (dt(2016, 1, 2), 1),
+                                                 (dt(2016, 1, 3), 1)],
+                                                names=['date', 'id'])
+                   )
+
+    # OPEN - CLOSED
+    assert_frame_equal(c.filter(df, DateRange(None, dt(2016, 1, 3))), df)
+    # CLOSED - OPEN
+    assert_frame_equal(c.filter(df, DateRange(dt(2016, 1, 1), None)), df)
+    # OPEN - OPEN
+    assert_frame_equal(c.filter(df, DateRange(None, None)), df)
+    # CLOSED - OPEN (far before data range)
+    assert_frame_equal(c.filter(df, DateRange(dt(2000, 1, 1), None)), df)
+    # CLOSED - OPEN (far after range)
+    assert(c.filter(df, DateRange(dt(2020, 1, 2), None)).empty)
+    # OPEN - CLOSED
+    assert_frame_equal(c.filter(df, DateRange(None, dt(2020, 1, 1))), df)
+    # CLOSED - CLOSED (after range)
+    assert(c.filter(df, DateRange(dt(2017, 1, 1), dt(2018, 1, 1))).empty)