Repository: pyjanitor-devs/pyjanitor
Pull request #1043 (2022-03-14 15:55:01)

[TST] Improve tests for `add_column` and `case_when`
# PR Description

Please describe the changes proposed in the pull request:
- Optimize placement of `df.copy()` till after the data checks in `add_column`
- Improve tests by specifying a "match" for the error message (otherwise ValueError is too generic). And not using `hypothesis` whenever possible.
- Remove some redundant /duplicated tests.



ericmjl: <!-- NETLIFY DEPLOY COMMENT GENERATED BY ACTIONS_NETLIFY -->
🚀 Deployed on https://deploy-preview-1043--pyjanitor.netlify.app
codecov[bot]: # [Codecov](https://codecov.io/gh/pyjanitor-devs/pyjanitor/pull/1043?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs) Report
> Merging [#1043](https://codecov.io/gh/pyjanitor-devs/pyjanitor/pull/1043?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs) (9e22457) into [dev](https://codecov.io/gh/pyjanitor-devs/pyjanitor/commit/f9dfe50e189608c5c265a6a754fc6c668d65c5d5?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyjanitor-devs) (f9dfe50) will **decrease** coverage by `0.00%`.
> The diff coverage is `100.00%`.


```diff
@@            Coverage Diff             @@
##              dev    #1043      +/-   ##
==========================================
- Coverage   97.74%   97.74%   -0.01%     
==========================================
  Files          77       77              
  Lines        2884     2883       -1     
==========================================
- Hits         2819     2818       -1     
  Misses         65       65              
```

ericmjl: @thatlittleboy I took the liberty to update the PR branch with the latest changes in `origin/dev`. We'll let the bots run and see what we get.
ericmjl: I vote we just merge.

Patch:
diff --git a/janitor/functions/add_columns.py b/janitor/functions/add_columns.py
index e326763..6055d24 100644
--- a/janitor/functions/add_columns.py
+++ b/janitor/functions/add_columns.py
@@ -64,14 +64,13 @@ def add_column(
         the number of rows in the DataFrame, repeat the list or tuple
         (R-style) to the end of the DataFrame.
     :returns: A pandas DataFrame with an added column.
-    :raises ValueError: if attempting to add a column that already exists.
-    :raises ValueError: if `value` has more elements that number of
+    :raises ValueError: If attempting to add a column that already exists.
+    :raises ValueError: If `value` has more elements that number of
         rows in the DataFrame.
-    :raises ValueError: if attempting to add an iterable of values with
+    :raises ValueError: If attempting to add an iterable of values with
         a length not equal to the number of DataFrame rows.
-    :raises ValueError: if `value` has length of `0`.
+    :raises ValueError: If `value` has length of `0`.
     """
-    df = df.copy()
     check("column_name", column_name, [str])
 
     if column_name in df.columns:
@@ -79,41 +78,39 @@ def add_column(
             f"Attempted to add column that already exists: " f"{column_name}."
         )
 
-    nrows = df.shape[0]
+    nrows = len(df)
 
     if hasattr(value, "__len__") and not isinstance(
         value, (str, bytes, bytearray)
     ):
+        len_value = len(value)
+
         # if `value` is a list, ndarray, etc.
-        if len(value) > nrows:
+        if len_value > nrows:
             raise ValueError(
                 "`value` has more elements than number of rows "
-                f"in your `DataFrame`. vals: {len(value)}, "
+                f"in your `DataFrame`. vals: {len_value}, "
                 f"df: {nrows}"
             )
-        if len(value) != nrows and not fill_remaining:
+        if len_value != nrows and not fill_remaining:
             raise ValueError(
                 "Attempted to add iterable of values with length"
                 " not equal to number of DataFrame rows"
             )
-
-        if len(value) == 0:
+        if not len_value:
             raise ValueError(
                 "`value` has to be an iterable of minimum length 1"
             )
-        len_value = len(value)
+
     elif fill_remaining:
         # relevant if a scalar val was passed, yet fill_remaining == True
         len_value = 1
         value = [value]
 
-    nrows = df.shape[0]
-
+    df = df.copy()
     if fill_remaining:
         times_to_loop = int(np.ceil(nrows / len_value))
-
         fill_values = list(value) * times_to_loop
-
         df[column_name] = fill_values[:nrows]
     else:
         df[column_name] = value
@@ -123,14 +120,19 @@ def add_column(
 
 @pf.register_dataframe_method
 def add_columns(
-    df: pd.DataFrame, fill_remaining: bool = False, **kwargs
+    df: pd.DataFrame,
+    fill_remaining: bool = False,
+    **kwargs,
 ) -> pd.DataFrame:
     """Add multiple columns to the dataframe.
 
     This method does not mutate the original DataFrame.
 
-    Method to augment `add_column` with ability to add multiple columns in
-    one go. This replaces the need for multiple `add_column` calls.
+    Method to augment
+    [`add_column`][janitor.functions.add_columns.add_column]
+    with ability to add multiple columns in
+    one go. This replaces the need for multiple
+    [`add_column`][janitor.functions.add_columns.add_column] calls.
 
     Usage is through supplying kwargs where the key is the col name and the
     values correspond to the values of the new DataFrame column.
@@ -152,7 +154,7 @@ def add_columns(
     :param fill_remaining: If value is a tuple or list that is smaller than
         the number of rows in the DataFrame, repeat the list or tuple
         (R-style) to the end of the DataFrame. (Passed to `add_column`)
-    :param kwargs: column, value pairs which are looped through in
+    :param **kwargs: Column, value pairs which are looped through in
         `add_column` calls.
     :returns: A pandas DataFrame with added columns.
     """
diff --git a/janitor/functions/case_when.py b/janitor/functions/case_when.py
index 1012229..37a372f 100644
--- a/janitor/functions/case_when.py
+++ b/janitor/functions/case_when.py
@@ -106,7 +106,7 @@ def case_when(df: pd.DataFrame, *args, column_name: str) -> pd.DataFrame:
         The 1-D array should be the same length as the DataFrame.
     :param column_name: Name of column to assign results to. A new column
         is created, if it does not already exist in the DataFrame.
-    :raises ValueError: if the condition fails to evaluate.
+    :raises ValueError: If the condition fails to evaluate.
     :returns: A pandas DataFrame.
     """
     conditions, targets, default = _case_when_checks(df, args, column_name)
@@ -125,11 +125,8 @@ def case_when(df: pd.DataFrame, *args, column_name: str) -> pd.DataFrame:
         # https://stackoverflow.com/a/46091127/7175713
         except Exception as e:
             raise ValueError(
-                f"""
-                condition{index} and value{index}
-                failed to evaluate.
-                Original error message: {e}
-                """
+                f"condition{index} and value{index} failed to evaluate. "
+                f"Original error message: {e}"
             ) from e
 
     return df.assign(**{column_name: default})
@@ -141,17 +138,12 @@ def _case_when_checks(df: pd.DataFrame, args, column_name):
     """
     if len(args) < 3:
         raise ValueError(
-            """
-            At least three arguments are required
-            for the `args` parameter.
-            """
+            "At least three arguments are required for the `args` parameter."
         )
     if len(args) % 2 != 1:
         raise ValueError(
-            """
-            It seems the `default` argument is missing
-            from the variable `args` parameter.
-            """
+            "It seems the `default` argument is missing from the variable "
+            "`args` parameter."
         )
 
     check("column_name", column_name, [str])
@@ -193,20 +185,14 @@ def _case_when_checks(df: pd.DataFrame, args, column_name):
         arr_ndim = default.ndim
     if arr_ndim != 1:
         raise ValueError(
-            """
-            The `default` argument should either be a 1-D array,
-            a scalar, or a callable that can evaluate to
-            a 1-D array.
-            """
+            "The `default` argument should either be a 1-D array, a scalar, "
+            "or a callable that can evaluate to a 1-D array."
         )
     if not isinstance(default, pd.Series):
         default = pd.Series(default)
     if default.size != len(df):
         raise ValueError(
-            """
-            The length of the `default` argument
-            should be equal to the length
-            of the DataFrame.
-            """
+            "The length of the `default` argument should be equal to the "
+            "length of the DataFrame."
         )
     return conditions, targets, default


Test patch:
diff --git a/tests/functions/test_add_column.py b/tests/functions/test_add_column.py
index 8df69fd..c5ca01c 100644
--- a/tests/functions/test_add_column.py
+++ b/tests/functions/test_add_column.py
@@ -11,30 +11,46 @@ from janitor.testing_utils.strategies import df_strategy
 @pytest.mark.functions
 @given(df=df_strategy())
 def test_add_column_add_integer(df):
-    # col_name wasn't a string
+    """col_name wasn't a string"""
     with pytest.raises(TypeError):
         df.add_column(column_name=42, value=42)
 
 
 @pytest.mark.functions
-@given(df=df_strategy())
-def test_add_column_already_exists(df):
-    # column already exists
-    with pytest.raises(ValueError):
-        df.add_column("a", 42)
+def test_add_column_already_exists(dataframe):
+    """column already exists"""
+    with pytest.raises(
+        ValueError,
+        match="Attempted to add column that already exists",
+    ):
+        dataframe.add_column("a", 42)
 
 
 @pytest.mark.functions
-@given(df=df_strategy())
-def test_add_column_too_many(df):
-    # too many values for dataframe num rows:
-    with pytest.raises(ValueError):
-        df.add_column("toomany", np.ones(100))
+def test_add_column_too_many(dataframe):
+    """too many values for dataframe num rows"""
+    with pytest.raises(
+        ValueError,
+        match="`value` has more elements than number of rows",
+    ):
+        dataframe.add_column("toomany", np.ones(100))
+
+
+@pytest.mark.functions
+def test_add_column_too_few_but_no_fill_remaining(dataframe):
+    """too few values for dataframe num rows"""
+    with pytest.raises(
+        ValueError,
+        match="add iterable of values with length not equal",
+    ):
+        dataframe.add_column("toomany", np.ones(2), fill_remaining=False)
 
 
 @pytest.mark.functions
 @given(df=df_strategy())
 def test_add_column_scalar(df):
+    """Checks `add_column` works as expected when adding a numeric scalar
+    to the column"""
     # column appears in DataFrame
     df = df.add_column("fortytwo", 42)
     assert "fortytwo" in df.columns
@@ -44,15 +60,18 @@ def test_add_column_scalar(df):
     series.name = "fortytwo"
     assert_series_equal(df["fortytwo"], series)
 
-    # scalar values are correct for strings
-    # also, verify sanity check excludes strings, which have a length:
-
 
 @pytest.mark.functions
 @given(df=df_strategy())
 def test_add_column_string(df):
+    """Checks `add_column` works as expected when adding a string scalar
+    to the column.
+
+    And ensure no errors are raised. The error checks on iterables should
+    exclude strings, which also have a length.
+    """
     df = df.add_column("fortythousand", "test string")
-    series = pd.Series(["test string"] * len(df))
+    series = pd.Series(["test string" for _ in range(len(df))])
     series.name = "fortythousand"
     assert_series_equal(df["fortythousand"], series)
 
@@ -65,35 +84,41 @@ def test_add_column_string(df):
 
 
 @pytest.mark.functions
-@given(df=df_strategy(), vals=st.lists(elements=st.integers()))
-def test_add_column_fill_remaining_iterable(df, vals):
-    if len(vals) > len(df) or len(vals) == 0:
-        with pytest.raises(ValueError):
-            df = df.add_column("fill_in_iterable", vals, fill_remaining=True)
-    else:
-        df = df.add_column("fill_in_iterable", vals, fill_remaining=True)
-        assert not pd.isna(df["fill_in_iterable"]).any()
+def test_add_column_iterator_repeat_subtraction(dataframe):
+    """Checks `add_column` works as expected when adding a pd Series
+    to the column"""
+    df = dataframe.add_column("city_pop", dataframe.a - dataframe.a)
+    assert df.city_pop.sum() == 0
+    assert df.city_pop.iloc[0] == 0
 
 
 @pytest.mark.functions
 @given(df=df_strategy())
 def test_add_column_fill_scalar(df):
-    # fill_remaining works - value is scalar
-    vals = 42
-    df = df.add_column("fill_in_scalar", vals, fill_remaining=True)
-    series = pd.Series([42] * len(df))
+    """Checks the `fill_remaining` parameter works as expected when value
+    is a scalar."""
+    df = df.add_column("fill_in_scalar", 42, fill_remaining=True)
+    series = pd.Series([42 for _ in range(len(df))])
     series.name = "fill_in_scalar"
     assert_series_equal(df["fill_in_scalar"], series)
 
 
 @pytest.mark.functions
-def test_add_column_single_value(dataframe):
-    df = dataframe.add_column("city_pop", 100)
-    assert df.city_pop.mean() == 100
+@given(df=df_strategy(), vals=st.lists(elements=st.integers()))
+def test_add_column_fill_remaining_iterable(df, vals: list):
+    """Checks the `fill_remaining` parameter works as expected."""
+    if len(vals) > len(df) or not vals:
+        with pytest.raises(ValueError):
+            df = df.add_column("fill_in_iterable", vals, fill_remaining=True)
+    else:
+        df = df.add_column("fill_in_iterable", vals, fill_remaining=True)
+        assert not pd.isna(df["fill_in_iterable"]).any()
 
 
 @pytest.mark.functions
 def test_add_column_iterator_repeat(dataframe):
+    """Fill remaining using a small iterator, with `fill_remaining` set to
+    True."""
     df = dataframe.add_column("city_pop", range(3), fill_remaining=True)
     assert df.city_pop.iloc[0] == 0
     assert df.city_pop.iloc[1] == 1
@@ -101,23 +126,3 @@ def test_add_column_iterator_repeat(dataframe):
     assert df.city_pop.iloc[3] == 0
     assert df.city_pop.iloc[4] == 1
     assert df.city_pop.iloc[5] == 2
-
-
-@pytest.mark.functions
-def test_add_column_raise_error(dataframe):
-    with pytest.raises(Exception):
-        dataframe.add_column("cities", 1)
-
-
-@pytest.mark.functions
-def test_add_column_iterator_repeat_subtraction(dataframe):
-    df = dataframe.add_column("city_pop", dataframe.a - dataframe.a)
-    assert df.city_pop.sum() == 0
-    assert df.city_pop.iloc[0] == 0
-
-
-@pytest.mark.functions
-@given(df=df_strategy())
-def test_add_column_checkequality(df):
-    new_df = df.add_column("fortytwo", 42)
-    assert new_df is not df
diff --git a/tests/functions/test_case_when.py b/tests/functions/test_case_when.py
index c38743f..b2f156d 100644
--- a/tests/functions/test_case_when.py
+++ b/tests/functions/test_case_when.py
@@ -1,11 +1,12 @@
+import numpy as np
 import pandas as pd
 import pytest
-import numpy as np
-from pandas.testing import assert_frame_equal
 from hypothesis import assume, given
+from pandas.testing import assert_frame_equal
+
 from janitor.testing_utils.strategies import (
-    df_strategy,
     categoricaldf_strategy,
+    df_strategy,
 )
 
 
@@ -38,27 +39,33 @@ def test_case_when_1():
     assert_frame_equal(result, expected)
 
 
-@given(df=df_strategy())
-def test_len_args(df):
+def test_len_args(dataframe):
     """Raise ValueError if `args` length is less than 3."""
-    with pytest.raises(ValueError):
-        df.case_when(df.a < 10, "less_than_10", column_name="a")
+    with pytest.raises(ValueError, match="three arguments are required"):
+        dataframe.case_when(dataframe.a < 10, "less_than_10", column_name="a")
 
 
-@given(df=df_strategy())
-def test_args_even(df):
+def test_args_even(dataframe):
     """Raise ValueError if `args` length is even."""
-    with pytest.raises(ValueError):
-        df.case_when(
-            df.a < 10, "less_than_10", df.a == 5, "five", column_name="a"
+    with pytest.raises(ValueError, match="`default` argument is missing"):
+        dataframe.case_when(
+            dataframe.a < 10,
+            "less_than_10",
+            dataframe.a == 5,
+            "five",
+            column_name="a",
         )
 
 
-@given(df=df_strategy())
-def test_column_name(df):
+def test_column_name(dataframe):
     """Raise TypeError if `column_name` is not a string."""
     with pytest.raises(TypeError):
-        df.case_when(df.a < 10, "less_than_10", df.a, column_name=("a",))
+        dataframe.case_when(
+            dataframe.a < 10,
+            "less_than_10",
+            dataframe.a,
+            column_name=("a",),
+        )
 
 
 @given(df=df_strategy())
@@ -72,9 +79,18 @@ def test_default_ndim(df):
 def test_default_length(df):
     """Raise ValueError if `default` length != len(df)."""
     assume(len(df) > 10)
-    with pytest.raises(ValueError):
+    with pytest.raises(
+        ValueError,
+        match=(
+            "length of the `default` argument should be equal to the length of"
+            " the DataFrame"
+        ),
+    ):
         df.case_when(
-            df.a < 10, "less_than_10", df.loc[:5, "a"], column_name="a"
+            df.a < 10,
+            "less_than_10",
+            df.loc[:5, "a"],
+            column_name="a",
         )