diff --git a/tests/lottery_predictor/test_analyze.py b/tests/lottery_predictor/test_analyze.py index 862fa80..e1a761f 100644 --- a/tests/lottery_predictor/test_analyze.py +++ b/tests/lottery_predictor/test_analyze.py @@ -1,3 +1,4 @@ +from datetime import datetime import pytest import pandas as pd @@ -5,8 +6,13 @@ import lottery_predictor.analyze as analyze def test_load_dataframe() -> None: - assert int(analyze.load_dataframe(table_name='PowerballDraw').count()['draw_date']) >= 1324 - assert int(analyze.load_dataframe(table_name='MegaMillionsDraw').count()['draw_date']) >= 95 + pb_start = datetime(2015, 10, 4) + mm_start = datetime(2025, 4, 5) + date_end = datetime(2026, 3, 8) + assert int(analyze.load_dataframe( + table_name='PowerballDraw', start_date=pb_start, end_date=date_end).count()['draw_date']) == 1324 + assert int(analyze.load_dataframe( + table_name='MegaMillionsDraw', start_date=mm_start, end_date=date_end).count()['draw_date']) == 96 with pytest.raises(ValueError): analyze.load_dataframe(table_name='SomNonExistentTableName')