test: make tests more flexible

Author: tjbell

tests/test_module1.py

@@ -53,7 +53,26 @@ def test_rename_columns_module1():
 
 @pytest.mark.test_concatenate_identifier_columns_module1
 def test_concatenate_identifier_columns_module1():
-    assert 'games:pd:concat:games:identifiers:axis:1:sort:False' in get_assignments(data), 'Concatenate the `games` and `identifiers` DataFrames.'
+    concat = False
+    frames = False
+    axis = False
+    sort = False
+
+    for string in get_assignments(data):
+        if 'games:pd:concat' in string:
+            concat = True
+        if 'games:identifiers' in string:
+            frames = True
+        if 'axis:1' in string:
+            axis = True
+        if 'sort:False' in string:
+            sort = True
+
+    assert concat, 'Are you calling `pd.concat()`?'
+    assert frames, 'Does the call to `pd.concat()` have a list of DataFrames to concatenate as the first argument? Make sure the frames are in the correct order.'
+    assert axis, 'Does the call to `pd.concat()` have a keyword argument of `axis` set to `1`?'
+    assert sort, 'Does the call to `pd.concat()` have a keyword argument of `sort` set to `False`?'
+
 
 @pytest.mark.test_fill_nan_values_module1
 def test_fill_nan_values_module1():

tests/test_module2.py

@@ -1,5 +1,7 @@
 import pytest
 import matplotlib
+import numpy as np
+import pandas as pd
 matplotlib.use('Agg')
 
 from .utils import get_assignments, get_calls
@@ -27,8 +29,10 @@ def test_select_attendance_module2():
         if 'multi3' not in attendance.attendance.columns:
             local_attendance.columns = ['year', 'attendance']
 
-        assert attendance.attendance.equals(local_attendance), 'Have you selected the attendance rows with `loc[]`?'
+        if attendance.attendance['attendance'].dtype == np.int64:
+            local_attendance.loc[:, 'attendance'] = pd.to_numeric(local_attendance.loc[:, 'attendance'])
 
+        assert attendance.attendance.equals(local_attendance), 'Have you selected the attendance rows with `loc[]`?'
     except ImportError:
         print('It looks as if `data.py` is incomplete.')
 
@@ -42,7 +46,28 @@ def test_convert_to_numeric_module2():
 
 @pytest.mark.test_plot_dataframe_module2
 def test_plot_dataframe_module2():
-    assert 'attendance:plot:x:year:y:attendance:figsize:15:7:kind:bar' in get_calls(attendance), 'Plot the `year` on the x-axis and the `attendance` on the y-axis of a bar plot. Adjust the size of the plot.'
+    plot = False
+    x = False
+    y = False
+    figsize = False
+    kind = False
+    for string in get_calls(attendance):
+        if 'attendance:plot' in string:
+            plot = True
+        if 'x:year' in string:
+            x = True
+        if 'y:attendance' in string:
+            y = True
+        if 'figsize:15:7' in string:
+            figsize = True
+        if 'kind:bar' in string:
+            kind = True
+
+    assert plot, 'Are you calling `plot()` on the `attendance` DataFrame?'
+    assert x, 'Does the call to `plot()` have a keyword argument of `x` set to `\'year\'`?'
+    assert y, 'Does the call to `plot()` have a keyword argument of `y` set to `\'attendance\'`?'
+    assert figsize, 'Does the call to `plot()` have a keyword argument of `figsize` set to `(15, 7)`?'
+    assert kind, 'Does the call to `plot()` have a keyword argument of `kind` set to `\'bar\'`?'
     assert 'plt:show' in get_calls(attendance), 'Have you shown the plot?'
 
 @pytest.mark.test_axis_labels_module2
@@ -52,4 +77,25 @@ def test_axis_labels_module2():
 
 @pytest.mark.test_mean_line_module2
 def test_mean_line_module2():
-    assert 'plt:axhline:y:attendance:attendance:mean:label:Mean:linestyle:--:color:green' in get_calls(attendance), 'Plot a green dashed line at the mean.'
+    axhline = False
+    y = False
+    label = False
+    linestyle = False
+    color = False
+    for string in get_calls(attendance):
+        if 'plt:axhline' in string:
+            axhline = True
+        if 'y:attendance:attendance:mean' in string:
+            y = True
+        if 'label:Mean' in string:
+            label = True
+        if 'linestyle:--' in string:
+            linestyle = True
+        if 'color:green' in string:
+            color = True
+
+    assert axhline, 'Are you calling `plt.axhline()`?'
+    assert y, 'Does the call to `plt.axhline()` have a keyword argument of `y` set to `attendance[\'attendance\'].mean()`?'
+    assert label, 'Does the call to `plt.axhline()` have a keyword argument of `label` set to `\'Mean\'`?'
+    assert linestyle, 'Does the call to `plt.axhline()` have a keyword argument of `linestyle` set to `\'--\'`?'
+    assert color, 'Does the call to `plt.axhline()` have a keyword argument of `color` set to `\'green\'`?'

tests/test_module3.py

@@ -8,7 +8,7 @@
 @pytest.mark.test_select_all_plays_module3
 def test_select_all_plays_module3():
     assert 'games' in dir(pitching), 'Have you imported `games` from `data`?'
-    assert 'plays:games:games:type:play' in get_assignments(pitching), 'Are you selecting just the rows that have a `type` of `play`?'
+    assert 'plays:games:games:type:play' in get_assignments(pitching), 'Are you selecting just the rows that have a `type` of `play`? Make sure you are using the shortcut method of selection.'
 
 @pytest.mark.test_select_all_strike_outs_module3
 def test_select_all_strike_outs_module3():
@@ -28,5 +28,26 @@ def test_apply_an_operation_to_multiple_columns_module3():
 
 @pytest.mark.test_change_plot_formatting_module3
 def test_change_plot_formatting_module3():
-    assert 'strike_outs:plot:x:year:y:strike_outs:kind:scatter:legend:Strike Outs' in get_calls(pitching), 'Create a scatter plot with the `year` as the x-axis and the number of `strikes_outs` on the y-axis.'
+    plot = False
+    x = False
+    y = False
+    kind = False
+    legend = False
+    for string in get_calls(pitching):
+        if 'strike_outs:plot' in string:
+            plot = True
+        if 'x:year' in string:
+            x = True
+        if 'y:strike_outs' in string:
+            y = True
+        if 'kind:scatter' in string:
+            kind = True
+        if 'legend:Strike Outs' in string:
+            legend = True
+
+    assert plot, 'Are you calling `plot()` on the `strike_outs` DataFrame?'
+    assert x, 'Does the call to `plot()` have a keyword argument of `x` set to `\'year\'`?'
+    assert y, 'Does the call to `plot()` have a keyword argument of `y` set to `\'strike_outs\'`?'
+    assert kind, 'Does the call to `plot()` have a keyword argument of `kind` set to `\'scatter\'`?'
+    assert legend, 'Have you chained a call to `legend()`?'
     assert 'plt:show' in get_calls(pitching), 'Show the scatter plot.'

tests/test_module4.py

@@ -45,7 +45,25 @@ def test_sort_values_module4():
 
 @pytest.mark.test_reshape_with_pivot_module4
 def test_reshape_with_pivot_module4():
-    assert 'hits:hits:pivot:index:inning:columns:hit_type:values:count' in get_assignments(offense), 'Make sure to reshape the `hits` DataFrame with the `pivot()` function. Ensure you have set the correct keyword arguments.'
+    pivot = False
+    index = False
+    columns = False
+    values = False
+
+    for string in get_assignments(offense):
+        if 'hits:hits:pivot' in string:
+            pivot = True
+        if 'index:inning' in string:
+            index = True
+        if 'columns:hit_type' in string:
+            columns = True
+        if 'values:count' in string:
+            values = True
+
+    assert pivot, 'Are you calling `pivot()` on the `hits` DataFrame?'
+    assert index, 'Does the call to `pivot()` have a keyword argument of `index` set to `\'inning\'`?'
+    assert columns, 'Does the call to `pivot()` have a keyword argument of `columns` set to `\'strike_outs\'`?'
+    assert values, 'Does the call to `pivot()` have a keyword argument of `values` set to `\'count\'`?'
 
 @pytest.mark.test_stacked_bar_plot_module4
 def test_stacked_bar_plot_module4():

tests/test_module5.py

@@ -54,7 +54,25 @@ def test_manage_column_labels_module5():
 
 @pytest.mark.test_merge_plate_appearances_module5
 def test_merge_plate_appearances_module5():
-    assert 'events_plus_pa:pd:merge:events:pa:how:outer:left_on:year:game_id:team:right_on:year:game_id:team' in get_assignments(defense), 'Have the `events` DataFrame and the `pa` DataFrame been merged with the correct keyword arguments?'
+    merge = False
+    how = False
+    left_on = False
+    right_on = False
+
+    for string in get_assignments(defense):
+        if 'events_plus_pa:pd:merge:events:pa' in string:
+            merge = True
+        if 'how:outer' in string:
+            how = True
+        if 'left_on:year:game_id:team' in string:
+            left_on = True
+        if 'right_on:year:game_id:team' in string:
+            right_on = True
+
+    assert merge, 'Are you calling `pd.merge()` to merge the `events` and `pa` DataFrames?'
+    assert how, 'Does the call to `pd.merge()` have a keyword argument of `how` set to `\'outer\'`?'
+    assert left_on, "Does the call to `pd.merge()` have a keyword argument of `left_on` set to `\'['year', 'game_id', 'team']\'`?"
+    assert right_on, "Does the call to `pd.merge()` have a keyword argument of `right_on` set to `\'['year', 'game_id', 'team']\'`?"
 
 @pytest.mark.test_merge_team_module5
 def test_merge_team_module5():
@@ -73,7 +91,26 @@ def test_calculate_der_module5():
 @pytest.mark.test_reshape_with_pivot_module5
 def test_reshape_with_pivot_module5():
     assert 'der:Name:defense:Name:loc:Attribute:defense:Name:year:Str:Index:Subscript:GtE:1978:Num:Compare:year:Str:defense:Str:DER:Str:List:Tuple:Index:Subscript:Assign' in get_assignments(defense, include_type=True), 'Select just the rows of the `defense` DataFrame with a year greater than 1978.'
-    assert 'der:der:pivot:index:year:columns:defense:values:DER' in get_assignments(defense), 'Reshape the `defense` DataFrame with the `pivot()` function and the correct keyword arguments.'
+
+    pivot = False
+    index = False
+    columns = False
+    values = False
+
+    for string in get_assignments(defense):
+        if 'der:der:pivot' in string:
+            pivot = True
+        if 'index:year' in string:
+            index = True
+        if 'columns:defense' in string:
+            columns = True
+        if 'values:DER' in string:
+            values = True
+
+    assert pivot, 'Are you calling `pivot()` on the `der` DataFrame?'
+    assert index, 'Does the call to `pivot()` have a keyword argument of `index` set to `\'year\'`?'
+    assert columns, 'Does the call to `pivot()` have a keyword argument of `columns` set to `\'defense\'`?'
+    assert values, 'Does the call to `pivot()` have a keyword argument of `values` set to `\'DER\'`?'
 
 @pytest.mark.test_plot_formatting_xticks_module5
 def test_plot_formatting_xticks_module5():