diff --git a/kenpompy/team.py b/kenpompy/team.py index 24cc624..2a96cd8 100644 --- a/kenpompy/team.py +++ b/kenpompy/team.py @@ -88,11 +88,14 @@ def get_schedule(browser, team=None, season=None): # Dataframe Tidying schedule_df = schedule_df[0] + # Teams 2010 and earlier do not show their team rank, add column for consistency + if(len(schedule_df.columns) == 10): + schedule_df.insert(1, 'Team Rank', '') schedule_df.columns = ['Date', 'Team Rank', 'Opponent Rank', 'Opponent Name', 'Result', 'Possession Number', 'A', 'Location', 'Record', 'Conference', 'B'] schedule_df = schedule_df.drop(columns = ['A', 'B']) schedule_df = schedule_df.fillna('') - schedule_df = schedule_df[schedule_df['Date'] != schedule_df['Team Rank']] + schedule_df = schedule_df[schedule_df['Date'] != schedule_df['Result']] schedule_df = schedule_df[schedule_df['Date'] != 'Date'] return schedule_df \ No newline at end of file diff --git a/tests/test_team.py b/tests/test_team.py index f721559..73b09b5 100644 --- a/tests/test_team.py +++ b/tests/test_team.py @@ -56,6 +56,14 @@ def test_get_schedule(browser): with pytest.raises(ValueError): kpteam.get_schedule(browser, team='Incorrect Team Name', season=2017) + centenary_expected = ['Sat Nov 11', '', '172', 'TCU', 'L, 72-66', '76', 'Away', '0-1', ''] + centenary_df = kpteam.get_schedule(browser, team='Centenary', season=2007) + assert [str(i) for i in centenary_df[centenary_df.Date == 'Sat Nov 11'].iloc[0].to_list()] == centenary_expected + assert centenary_df.shape == (31, 9) + + with pytest.raises(ValueError): + kpteam.get_schedule(browser, team='Centenary', season=2017) + # Make sure that the valid team check is triggered with pytest.raises(ValueError): kpteam.get_schedule(browser, season = '2013', team="LMU") \ No newline at end of file