import pytest
import pandas as pd

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
    with pytest.raises(ValueError):
        analyze.load_dataframe(table_name='SomNonExistentTableName')


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

    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