diff --git a/tests/test_core.py b/tests/test_core.py index e217979..800449d 100644 --- a/tests/test_core.py +++ b/tests/test_core.py @@ -2,18 +2,20 @@ import datetime import random from typing import Literal, Mapping, Sequence +import pytest from fincal.core import AllFrequencies, Frequency, Series, TimeSeriesCore from fincal.fincal import create_date_series +from fincal.utils import FincalOptions class TestFrequency: def test_creation(self): - D = Frequency('daily', 'days', 1, 1, 'D') + D = Frequency("daily", "days", 1, 1, "D") assert D.days == 1 - assert D.symbol == 'D' - assert D.name == 'daily' + assert D.symbol == "D" + assert D.name == "daily" assert D.value == 1 - assert D.freq_type == 'days' + assert D.freq_type == "days" def create_test_data( @@ -60,9 +62,9 @@ def create_test_data( class TestAllFrequencies: def test_attributes(self): - assert hasattr(AllFrequencies, 'D') - assert hasattr(AllFrequencies, 'M') - assert hasattr(AllFrequencies, 'Q') + assert hasattr(AllFrequencies, "D") + assert hasattr(AllFrequencies, "M") + assert hasattr(AllFrequencies, "Q") def test_days(self): assert AllFrequencies.D.days == 1 @@ -70,64 +72,135 @@ class TestAllFrequencies: assert AllFrequencies.Q.days == 91 def test_symbol(self): - assert AllFrequencies.H.symbol == 'H' - assert AllFrequencies.W.symbol == 'W' + assert AllFrequencies.H.symbol == "H" + assert AllFrequencies.W.symbol == "W" def test_values(self): assert AllFrequencies.H.value == 6 assert AllFrequencies.Y.value == 1 def test_type(self): - assert AllFrequencies.Q.freq_type == 'months' - assert AllFrequencies.W.freq_type == 'days' + assert AllFrequencies.Q.freq_type == "months" + assert AllFrequencies.W.freq_type == "days" class TestSeries: def test_creation(self): - series = Series([1, 2, 3, 4, 5, 6, 7], data_type='number') + series = Series([1, 2, 3, 4, 5, 6, 7], data_type="number") assert series.dtype == float assert series[2] == 3 - dates = create_date_series('2021-01-01', '2021-01-31', frequency='D') - series = Series(dates, data_type='date') + dates = create_date_series("2021-01-01", "2021-01-31", frequency="D") + series = Series(dates, data_type="date") assert series.dtype == datetime.datetime class TestTimeSeriesCore: - data = [('2021-01-01', 220), ('2021-02-01', 230), ('2021-03-01', 240)] + data = [("2021-01-01", 220), ("2021-02-01", 230), ("2021-03-01", 240)] def test_repr_str(self): - ts = TimeSeriesCore(self.data, frequency='M') - assert str(ts) in repr(ts).replace('\t', ' ') + ts = TimeSeriesCore(self.data, frequency="M") + assert str(ts) in repr(ts).replace("\t", " ") data = create_test_data(frequency="D", eomonth=False, n=50, gaps=0, month_position="start", date_as_str=True) ts = TimeSeriesCore(data, frequency="D") - assert '...' in str(ts) - assert '...' in repr(ts) + assert "..." in str(ts) + assert "..." in repr(ts) def test_creation(self): - ts = TimeSeriesCore(self.data, frequency='M') + ts = TimeSeriesCore(self.data, frequency="M") assert isinstance(ts, TimeSeriesCore) assert isinstance(ts, Mapping) def test_getitem(self): - ts = TimeSeriesCore(self.data, frequency='M') + ts = TimeSeriesCore(self.data, frequency="M") assert ts.dates[0] == datetime.datetime(2021, 1, 1, 0, 0) assert ts.values[0] == 220 - assert ts['2021-01-01'][1] == 220 - assert len(ts[ts.dates > '2021-01-01']) == 2 - assert ts[ts.dates == '2021-02-01'].iloc[0][1] == 230 + assert ts["2021-01-01"][1] == 220 + assert len(ts[ts.dates > "2021-01-01"]) == 2 + assert ts[ts.dates == "2021-02-01"].iloc[0][1] == 230 assert ts.iloc[2][0] == datetime.datetime(2021, 3, 1) assert len(ts.iloc[:2]) == 2 + with pytest.raises(KeyError): + ts["2021-02-03"] + subset_ts = ts[["2021-01-01", "2021-03-01"]] + assert len(subset_ts) == 2 + assert isinstance(subset_ts, TimeSeriesCore) + assert subset_ts.iloc[1][1] == 240 + + def test_get(self): + ts = TimeSeriesCore(self.data, frequency="M") + assert ts.dates[0] == datetime.datetime(2021, 1, 1, 0, 0) + assert ts.values[0] == 220 + assert ts.get("2021-01-01")[1] == 220 + assert ts.get("2021-02-15") is None + assert ts.get("2021-02-23", -1) == -1 + assert ts.get("2021-02-10", closest="previous")[1] == 230 + assert ts.get("2021-02-10", closest="next")[1] == 240 + FincalOptions.get_closest = "previous" + assert ts.get("2021-02-10")[1] == 230 + FincalOptions.get_closest = "next" + assert ts.get("2021-02-10")[1] == 240 def test_contains(self): - ts = TimeSeriesCore(self.data, frequency='M') + ts = TimeSeriesCore(self.data, frequency="M") assert datetime.datetime(2021, 1, 1) in ts - assert '2021-01-01' in ts - assert '2021-01-14' not in ts + assert "2021-01-01" in ts + assert "2021-01-14" not in ts def test_items(self): - ts = TimeSeriesCore(self.data, frequency='M') + ts = TimeSeriesCore(self.data, frequency="M") for i, j in ts.items(): assert j == self.data[0][1] break + + def test_special_keys(self): + ts = TimeSeriesCore(self.data, frequency="M") + dates = ts["dates"] + values = ts["values"] + assert isinstance(dates, Series) + assert isinstance(values, Series) + assert len(dates) == 3 + assert len(values) == 3 + assert dates[0] == datetime.datetime(2021, 1, 1, 0, 0) + assert values[0] == 220 + + +class TestTimeSeriesCoreHeadTail: + data = [ + ("2021-01-01", 220), + ("2021-02-01", 230), + ("2021-03-01", 240), + ("2021-04-01", 250), + ("2021-05-01", 260), + ("2021-06-01", 270), + ("2021-07-01", 280), + ("2021-08-01", 290), + ("2021-09-01", 300), + ("2021-10-01", 310), + ("2021-11-01", 320), + ("2021-12-01", 330), + ] + + def test_head(self): + ts = TimeSeriesCore(self.data, frequency="M") + assert len(ts.head()) == 6 + assert len(ts.head(3)) == 3 + assert isinstance(ts.head(), TimeSeriesCore) + head_ts = ts.head(6) + assert head_ts.iloc[-1][1] == 270 + + def test_tail(self): + ts = TimeSeriesCore(self.data, frequency="M") + assert len(ts.tail()) == 6 + assert len(ts.tail(8)) == 8 + assert isinstance(ts.tail(), TimeSeriesCore) + tail_ts = ts.tail(6) + assert tail_ts.iloc[0][1] == 280 + + def test_head_tail(self): + ts = TimeSeriesCore(self.data, frequency="M") + head_tail_ts = ts.head(8).tail(2) + assert isinstance(head_tail_ts, TimeSeriesCore) + assert "2021-07-01" in head_tail_ts + assert head_tail_ts.iloc[1][1] == 290