diff --git a/tests/lottery_predictor/test_analyze.py b/tests/lottery_predictor/test_analyze.py index e1a761f..7f6d48f 100644 --- a/tests/lottery_predictor/test_analyze.py +++ b/tests/lottery_predictor/test_analyze.py @@ -18,14 +18,21 @@ def test_load_dataframe() -> None: def test_least_and_most_common_number() -> None: - pb_df = analyze.load_dataframe(table_name='PowerballDraw') - mm_df = analyze.load_dataframe(table_name='MegaMillionsDraw') - assert len(analyze.get_most_common_number(pb_df, top=5)) == 5 - assert len(analyze.get_least_common_number(pb_df, bottom=5)) == 5 - assert len(analyze.get_most_common_number(mm_df, top=5)) == 5 - assert len(analyze.get_least_common_number(mm_df, bottom=5)) == 5 + pb_start = datetime(2015, 10, 4) + mm_start = datetime(2025, 4, 5) + date_end = datetime(2026, 3, 8) - assert len(analyze.get_most_common_number(pb_df, columns=['powerball'], top=1)) == 1 - assert len(analyze.get_least_common_number(pb_df, columns=['powerball'], bottom=1)) == 1 - assert len(analyze.get_most_common_number(mm_df, columns=['mega_ball'], top=1)) == 1 - assert len(analyze.get_least_common_number(mm_df, columns=['mega_ball'], bottom=1)) == 1 + pb_df = analyze.load_dataframe( + table_name='PowerballDraw', start_date=pb_start, end_date=date_end) + mm_df = analyze.load_dataframe( + table_name='MegaMillionsDraw', start_date=mm_start, end_date=date_end) + + assert analyze.get_most_common_number(pb_df, top=5) == [61, 21, 28, 23, 33] + assert analyze.get_least_common_number(pb_df, bottom=5) == [13, 49, 46, 26, 34] + assert analyze.get_most_common_number(mm_df, top=5) == [42, 40, 18, 10, 49] + assert analyze.get_least_common_number(mm_df, bottom=5) == [35, 51, 3, 1, 20] + + assert analyze.get_most_common_number(pb_df, columns=['powerball'], top=1) == [4] + assert analyze.get_least_common_number(pb_df, columns=['powerball'], bottom=1) == [16] + assert analyze.get_most_common_number(mm_df, columns=['mega_ball'], top=1) == [1] + assert analyze.get_least_common_number(mm_df, columns=['mega_ball'], bottom=1) == [3]