Skip to content

Commit cb0265d

Browse files
committed
REFACTORED TESTS
1 parent 2532095 commit cb0265d

19 files changed

+3936
-2305
lines changed
Lines changed: 341 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,341 @@
1+
from datetime import datetime
2+
3+
import pandas as pd
4+
import pytest
5+
import pytz
6+
7+
from staircase import Stairs
8+
from staircase.constants import inf
9+
10+
11+
def pytest_generate_tests(metafunc):
12+
if "date_func" in metafunc.fixturenames:
13+
metafunc.parametrize(
14+
"date_func",
15+
["pandas", "pydatetime", "numpy", "pandas_tz", "pydatetime_tz"],
16+
indirect=True,
17+
)
18+
19+
20+
@pytest.fixture
21+
def date_func(request):
22+
# returns a func which takes a pandas timestamp
23+
if request.param == "pandas":
24+
return lambda x: x
25+
elif request.param == "pydatetime":
26+
return pd.Timestamp.to_pydatetime
27+
elif request.param == "numpy":
28+
return pd.Timestamp.to_datetime64
29+
elif request.param == "pandas_tz":
30+
return lambda ts: pd.Timestamp.tz_localize(
31+
ts, pytz.timezone("Australia/Sydney")
32+
)
33+
elif request.param == "pydatetime_tz":
34+
return lambda ts: (
35+
pd.Timestamp.tz_localize(
36+
ts, pytz.timezone("Australia/Sydney")
37+
).to_pydatetime()
38+
)
39+
else:
40+
assert False, "should not happen"
41+
42+
43+
def timestamp(*args, date_func, **kwargs):
44+
ts = pd.Timestamp(*args, **kwargs)
45+
return date_func(ts)
46+
47+
48+
def assert_expected_type(stairs, date_func):
49+
if stairs._data is None:
50+
return
51+
example_type = timestamp(2020, 1, 1, date_func=date_func)
52+
example_type = pd.Timestamp(
53+
example_type
54+
) # pandas natively converts datetimes to timestamps
55+
assert all(
56+
[type(example_type) == type(x) for x in stairs._data.index]
57+
), "Unexpected type in step points"
58+
if isinstance(example_type, (pd.Timestamp, datetime)):
59+
assert all(
60+
[example_type.tzinfo == x.tzinfo for x in stairs._data.index]
61+
), "Unexpected timezone in step points"
62+
63+
64+
def s1(date_func):
65+
int_seq1 = Stairs(initial_value=0)
66+
int_seq1.layer(
67+
timestamp(2020, 1, 1, date_func=date_func),
68+
timestamp(2020, 1, 10, date_func=date_func),
69+
2,
70+
)
71+
int_seq1.layer(
72+
timestamp(2020, 1, 3, date_func=date_func),
73+
timestamp(2020, 1, 5, date_func=date_func),
74+
2.5,
75+
)
76+
int_seq1.layer(
77+
timestamp(2020, 1, 6, date_func=date_func),
78+
timestamp(2020, 1, 7, date_func=date_func),
79+
-2.5,
80+
)
81+
int_seq1.layer(
82+
timestamp(2020, 1, 7, date_func=date_func),
83+
timestamp(2020, 1, 10, date_func=date_func),
84+
-2.5,
85+
)
86+
return int_seq1
87+
88+
89+
def s2(date_func):
90+
int_seq2 = Stairs(initial_value=0)
91+
int_seq2.layer(
92+
timestamp(2020, 1, 1, date_func=date_func),
93+
timestamp(2020, 1, 7, date_func=date_func),
94+
-2.5,
95+
)
96+
int_seq2.layer(
97+
timestamp(2020, 1, 8, date_func=date_func),
98+
timestamp(2020, 1, 10, date_func=date_func),
99+
5,
100+
)
101+
int_seq2.layer(
102+
timestamp(2020, 1, 2, date_func=date_func),
103+
timestamp(2020, 1, 5, date_func=date_func),
104+
4.5,
105+
)
106+
int_seq2.layer(
107+
timestamp(2020, 1, 2, 12, date_func=date_func),
108+
timestamp(2020, 1, 4, date_func=date_func),
109+
-2.5,
110+
)
111+
return int_seq2
112+
113+
114+
def s3(date_func): # boolean
115+
int_seq = Stairs(initial_value=0)
116+
int_seq.layer(
117+
timestamp(2020, 1, 10, date_func=date_func),
118+
timestamp(2020, 1, 30, date_func=date_func),
119+
1,
120+
)
121+
int_seq.layer(
122+
timestamp(2020, 1, 12, date_func=date_func),
123+
timestamp(2020, 1, 13, date_func=date_func),
124+
-1,
125+
)
126+
int_seq.layer(
127+
timestamp(2020, 1, 15, date_func=date_func),
128+
timestamp(2020, 1, 18, date_func=date_func),
129+
-1,
130+
)
131+
int_seq.layer(
132+
timestamp(2020, 1, 20, 12, date_func=date_func),
133+
timestamp(2020, 1, 21, date_func=date_func),
134+
-1,
135+
)
136+
int_seq.layer(
137+
timestamp(2020, 1, 23, date_func=date_func),
138+
timestamp(2020, 1, 23, 12, date_func=date_func),
139+
-1,
140+
)
141+
int_seq.layer(
142+
timestamp(2020, 1, 27, date_func=date_func),
143+
timestamp(2020, 1, 29, 12, date_func=date_func),
144+
-1,
145+
)
146+
return int_seq
147+
148+
149+
def s4(date_func): # boolean
150+
int_seq = Stairs(initial_value=0)
151+
int_seq.layer(
152+
timestamp(2020, 1, 9, date_func=date_func),
153+
timestamp(2020, 1, 29, date_func=date_func),
154+
1,
155+
)
156+
int_seq.layer(
157+
timestamp(2020, 1, 10, 12, date_func=date_func),
158+
timestamp(2020, 1, 12, date_func=date_func),
159+
-1,
160+
)
161+
int_seq.layer(
162+
timestamp(2020, 1, 12, 12, date_func=date_func),
163+
timestamp(2020, 1, 13, date_func=date_func),
164+
-1,
165+
)
166+
int_seq.layer(
167+
timestamp(2020, 1, 20, date_func=date_func),
168+
timestamp(2020, 1, 23, date_func=date_func),
169+
-1,
170+
)
171+
int_seq.layer(
172+
timestamp(2020, 1, 26, date_func=date_func),
173+
timestamp(2020, 1, 26, 12, date_func=date_func),
174+
-1,
175+
)
176+
int_seq.layer(
177+
timestamp(2020, 1, 27, date_func=date_func),
178+
timestamp(2020, 1, 28, 12, date_func=date_func),
179+
-1,
180+
)
181+
return int_seq
182+
183+
184+
@pytest.fixture
185+
def s1_fix():
186+
return s1()
187+
188+
189+
@pytest.fixture
190+
def s2_fix():
191+
return s2()
192+
193+
194+
@pytest.fixture
195+
def s3_fix():
196+
return s3()
197+
198+
199+
@pytest.fixture
200+
def s4_fix():
201+
return s4()
202+
203+
204+
def test_add_dates(date_func):
205+
expected_step_changes = pd.Series(
206+
{
207+
timestamp("2020-01-01 00:00:00", date_func=date_func): -0.5,
208+
timestamp("2020-01-02 00:00:00", date_func=date_func): 4.5,
209+
timestamp("2020-01-02 12:00:00", date_func=date_func): -2.5,
210+
timestamp("2020-01-03 00:00:00", date_func=date_func): 2.5,
211+
timestamp("2020-01-04 00:00:00", date_func=date_func): 2.5,
212+
timestamp("2020-01-05 00:00:00", date_func=date_func): -7.0,
213+
timestamp("2020-01-06 00:00:00", date_func=date_func): -2.5,
214+
timestamp("2020-01-07 00:00:00", date_func=date_func): 2.5,
215+
timestamp("2020-01-08 00:00:00", date_func=date_func): 5,
216+
timestamp("2020-01-10 00:00:00", date_func=date_func): -4.5,
217+
}
218+
)
219+
result = s1(date_func) + s2(date_func)
220+
pd.testing.assert_series_equal(
221+
result.step_changes,
222+
expected_step_changes,
223+
check_names=False,
224+
check_index_type=False,
225+
)
226+
assert_expected_type(result, date_func)
227+
228+
229+
def test_subtract_dates(date_func):
230+
expected_step_changes = pd.Series(
231+
{
232+
timestamp("2020-01-01 00:00:00", date_func=date_func): 4.5,
233+
timestamp("2020-01-02 00:00:00", date_func=date_func): -4.5,
234+
timestamp("2020-01-02 12:00:00", date_func=date_func): 2.5,
235+
timestamp("2020-01-03 00:00:00", date_func=date_func): 2.5,
236+
timestamp("2020-01-04 00:00:00", date_func=date_func): -2.5,
237+
timestamp("2020-01-05 00:00:00", date_func=date_func): 2.0,
238+
timestamp("2020-01-06 00:00:00", date_func=date_func): -2.5,
239+
timestamp("2020-01-07 00:00:00", date_func=date_func): -2.5,
240+
timestamp("2020-01-08 00:00:00", date_func=date_func): -5,
241+
timestamp("2020-01-10 00:00:00", date_func=date_func): 5.5,
242+
}
243+
)
244+
result = s1(date_func) - s2(date_func)
245+
pd.testing.assert_series_equal(
246+
result.step_changes,
247+
expected_step_changes,
248+
check_names=False,
249+
check_index_type=False,
250+
)
251+
assert_expected_type(result, date_func)
252+
253+
254+
def test_multiply_dates(date_func):
255+
expected_step_changes = pd.Series(
256+
{
257+
timestamp("2020-01-01 00:00:00", date_func=date_func): -5.0,
258+
timestamp("2020-01-02 00:00:00", date_func=date_func): 9.0,
259+
timestamp("2020-01-02 12:00:00", date_func=date_func): -5.0,
260+
timestamp("2020-01-03 00:00:00", date_func=date_func): -1.25,
261+
timestamp("2020-01-04 00:00:00", date_func=date_func): 11.25,
262+
timestamp("2020-01-05 00:00:00", date_func=date_func): -14.0,
263+
timestamp("2020-01-06 00:00:00", date_func=date_func): 6.25,
264+
timestamp("2020-01-07 00:00:00", date_func=date_func): -1.25,
265+
timestamp("2020-01-08 00:00:00", date_func=date_func): -2.5,
266+
timestamp("2020-01-10 00:00:00", date_func=date_func): 2.5,
267+
}
268+
)
269+
result = s1(date_func) * s2(date_func)
270+
pd.testing.assert_series_equal(
271+
result.step_changes,
272+
expected_step_changes,
273+
check_names=False,
274+
check_index_type=False,
275+
)
276+
assert_expected_type(result, date_func)
277+
278+
279+
def test_multiply_dates_scalar(date_func):
280+
expected_step_changes = pd.Series(
281+
{
282+
timestamp("2020-01-01 00:00:00", date_func=date_func): 6.0,
283+
timestamp("2020-01-03 00:00:00", date_func=date_func): 7.5,
284+
timestamp("2020-01-05 00:00:00", date_func=date_func): -7.5,
285+
timestamp("2020-01-06 00:00:00", date_func=date_func): -7.5,
286+
timestamp("2020-01-10 00:00:00", date_func=date_func): 1.5,
287+
}
288+
)
289+
result = s1(date_func) * 3
290+
pd.testing.assert_series_equal(
291+
result.step_changes,
292+
expected_step_changes,
293+
check_names=False,
294+
check_index_type=False,
295+
)
296+
assert_expected_type(result, date_func)
297+
298+
299+
def test_divide_dates(date_func):
300+
expected_step_changes = pd.Series(
301+
{
302+
timestamp("2020-01-01 00:00:00", date_func=date_func): -1.3333333333333333,
303+
timestamp("2020-01-02 00:00:00", date_func=date_func): 2.0,
304+
timestamp("2020-01-02 12:00:00", date_func=date_func): 3.3333333333333335,
305+
timestamp("2020-01-03 00:00:00", date_func=date_func): 5.0,
306+
timestamp("2020-01-04 00:00:00", date_func=date_func): -7.5,
307+
timestamp("2020-01-05 00:00:00", date_func=date_func): -2.833333333333333,
308+
timestamp("2020-01-06 00:00:00", date_func=date_func): 1.6666666666666665,
309+
timestamp("2020-01-07 00:00:00", date_func=date_func): -0.8333333333333333,
310+
timestamp("2020-01-08 00:00:00", date_func=date_func): 0.4166666666666667,
311+
timestamp("2020-01-10 00:00:00", date_func=date_func): 0.08333333333333333,
312+
}
313+
)
314+
result = s1(date_func) / (s2(date_func) + 1)
315+
pd.testing.assert_series_equal(
316+
result.step_changes,
317+
expected_step_changes,
318+
check_names=False,
319+
check_index_type=False,
320+
)
321+
assert_expected_type(result, date_func)
322+
323+
324+
def test_divide_dates_scalar(date_func):
325+
expected_step_changes = pd.Series(
326+
{
327+
timestamp("2020-01-01 00:00:00", date_func=date_func): 4.0,
328+
timestamp("2020-01-03 00:00:00", date_func=date_func): 5.0,
329+
timestamp("2020-01-05 00:00:00", date_func=date_func): -5.0,
330+
timestamp("2020-01-06 00:00:00", date_func=date_func): -5.0,
331+
timestamp("2020-01-10 00:00:00", date_func=date_func): 1.0,
332+
}
333+
)
334+
result = s1(date_func) / 0.5
335+
pd.testing.assert_series_equal(
336+
result.step_changes,
337+
expected_step_changes,
338+
check_names=False,
339+
check_index_type=False,
340+
)
341+
assert_expected_type(result, date_func)

0 commit comments

Comments
 (0)