Added tests for date_series & some for TimeSeries

tests/test_fincal.py

@@ -1,21 +1,151 @@
+import datetime
+import os
+import random
+from typing import Literal, Sequence
 
-import unittest
-
-from fincal.fincal import TimeSeries
-
-
-class TestFincal(unittest.TestCase):
-    def test_ts(self):
-        data = [
-            ('2020-01-01',  23),
-            ('2020-01-02',  24),
-            ('2020-01-03',  25),
-            ('2020-01-06',  26),
-            ('2020-01-07',  27),
-            ('2020-01-08',  28),
-            ('2020-01-10',  29),
-            ('2020-01-11',  30)
-        ]
-        time_series = TimeSeries(data)
-        time_series.ffill(inplace=True)
-        self.assertEqual(len(time_series.time_series), 11)
+import pytest
+from fincal.core import Frequency
+from fincal.fincal import TimeSeries, create_date_series
+
+THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+sample_data_path = os.path.join(THIS_DIR, "data")
+
+
+def create_test_data(
+    frequency: str,
+    eomonth: bool,
+    n: int,
+    gaps: float,
+    month_position: Literal["start", "middle", "end"],
+    date_as_str: bool,
+    as_outer_type: Literal['dict', 'list'] = 'list',
+    as_inner_type: Literal['dict', 'list', 'tuple'] = 'tuple'
+) -> Sequence[tuple]:
+    start_dates = {
+        "start": datetime.datetime(2016, 1, 1),
+        "middle": datetime.datetime(2016, 1, 15),
+        "end": datetime.datetime(2016, 1, 31),
+    }
+    end_date = datetime.datetime(2021, 12, 31)
+    dates = create_date_series(start_dates[month_position], end_date, frequency=frequency, eomonth=eomonth)
+    dates = dates[:n]
+    if gaps:
+        num_gaps = int(len(dates) * gaps)
+        to_remove = random.sample(dates, num_gaps)
+        for i in to_remove:
+            dates.remove(i)
+    if date_as_str:
+        dates = [i.strftime('%Y-%m-%d') for i in dates]
+
+    values = [random.randint(8000, 90000)/100 for _ in dates]
+
+    data = list(zip(dates, values))
+    if as_outer_type == 'list':
+        if as_inner_type == 'list':
+            data = [list(i) for i in data]
+        elif as_inner_type == 'dict[1]':
+            data = [dict((i,)) for i in data]
+        elif as_inner_type == 'dict[2]':
+            data = [dict(date=i, value=j) for i, j in data]
+    elif as_outer_type == 'dict':
+        data = dict(data)
+
+    return data
+
+
+class TestDateSeries:
+    def test_daily(self):
+        start_date = datetime.datetime(2020, 1, 1)
+        end_date = datetime.datetime(2020, 12, 31)
+        d = create_date_series(start_date, end_date, frequency="D")
+        assert len(d) == 366
+
+        start_date = datetime.datetime(2017, 1, 1)
+        end_date = datetime.datetime(2017, 12, 31)
+        d = create_date_series(start_date, end_date, frequency="D")
+        assert len(d) == 365
+
+        with pytest.raises(ValueError):
+            create_date_series(start_date, end_date, frequency="D", eomonth=True)
+
+    def test_monthly(self):
+        start_date = datetime.datetime(2020, 1, 1)
+        end_date = datetime.datetime(2020, 12, 31)
+        d = create_date_series(start_date, end_date, frequency="M")
+        assert len(d) == 12
+
+        d = create_date_series(start_date, end_date, frequency="M", eomonth=True)
+        assert datetime.datetime(2020, 2, 29) in d
+
+        start_date = datetime.datetime(2020, 1, 31)
+        d = create_date_series(start_date, end_date, frequency="M")
+        assert datetime.datetime(2020, 2, 29) in d
+        assert datetime.datetime(2020, 8, 31) in d
+        assert datetime.datetime(2020, 10, 30) not in d
+
+        start_date = datetime.datetime(2020, 2, 29)
+        d = create_date_series(start_date, end_date, frequency="M")
+        assert len(d) == 11
+        assert datetime.datetime(2020, 2, 29) in d
+        assert datetime.datetime(2020, 8, 31) not in d
+        assert datetime.datetime(2020, 10, 29) in d
+
+    def test_quarterly(self):
+        start_date = datetime.datetime(2018, 1, 1)
+        end_date = datetime.datetime(2020, 12, 31)
+        d = create_date_series(start_date, end_date, frequency="Q")
+        assert len(d) == 12
+
+        d = create_date_series(start_date, end_date, frequency="Q", eomonth=True)
+        assert datetime.datetime(2020, 4, 30) in d
+
+        start_date = datetime.datetime(2020, 1, 31)
+        d = create_date_series(start_date, end_date, frequency="Q")
+        assert len(d) == 4
+        assert datetime.datetime(2020, 2, 29) not in d
+        assert max(d) == datetime.datetime(2020, 10, 31)
+
+        start_date = datetime.datetime(2020, 2, 29)
+        d = create_date_series(start_date, end_date, frequency="Q")
+        assert datetime.datetime(2020, 2, 29) in d
+        assert datetime.datetime(2020, 8, 31) not in d
+        assert datetime.datetime(2020, 11, 29) in d
+
+        d = create_date_series(start_date, end_date, frequency="Q", eomonth=True)
+        assert datetime.datetime(2020, 11, 30) in d
+
+
+class TestFincal:
+    def test_creation(self):
+        data = create_test_data(frequency='D', eomonth=False, n=50, gaps=0, month_position='start', date_as_str=True)
+        time_series = TimeSeries(data, frequency="D")
+        assert len(time_series) == 50
+        assert isinstance(time_series.frequency, Frequency)
+        assert time_series.frequency.days == 1
+
+        ffill_data = time_series.ffill()
+        assert len(ffill_data) == 50
+
+        data = create_test_data(frequency='D', eomonth=False, n=500, gaps=0.1, month_position='start', date_as_str=True)
+        time_series = TimeSeries(data, frequency="D")
+        assert len(time_series) == 450
+
+    def test_ffill(self):
+        data = create_test_data(frequency='D', eomonth=False, n=500, gaps=0.1, month_position='start', date_as_str=True)
+        time_series = TimeSeries(data, frequency="D")
+        ffill_data = time_series.ffill()
+        assert len(ffill_data) > 498
+
+        ffill_data = time_series.ffill(inplace=True)
+        assert ffill_data is None
+        assert len(time_series) > 498
+
+    def test_slicing(self):
+        data = create_test_data(frequency='D', eomonth=False, n=50, gaps=0, month_position='start', date_as_str=True)
+        time_series = TimeSeries(data, frequency="D")
+        assert time_series[0] is not None
+        assert time_series[:3] is not None
+        assert time_series[5:7] is not None
+        assert isinstance(time_series[0], tuple)
+        assert isinstance(time_series[10:20], list)
+        assert len(time_series[10:20]) == 10