Fixed null_checker tests to find value errors and key errors

George Robertson · George Robertson · commit 74498789a483 · 2019-09-10T16:18:53.000+01:00
diff --git a/codonPython/check_null.py b/codonPython/check_null.py
@@ -1,7 +1,7 @@
 import numpy
 import pandas as pd
 
-def check_null(dataframe: pd.DataFrame, columns_to_be_checked: tuple) -> bool:
+def check_null(dataframe: pd.DataFrame, columns_to_be_checked: list) -> bool:
     """
     Checks a pandas dataframe for null values
 
@@ -11,7 +11,7 @@ def check_null(dataframe: pd.DataFrame, columns_to_be_checked: tuple) -> bool:
     ----------
     data : pandas.DataFrame
         Dataframe to read
-    columns_to_be_checked: tuple
+    columns_to_be_checked: list
         Given dataframe columns to be checked for null values
 
     Returns
@@ -21,21 +21,22 @@ def check_null(dataframe: pd.DataFrame, columns_to_be_checked: tuple) -> bool:
 
     Examples
     --------
-    >>> check_null(dataframe = pd.DataFrame({'col1': [1,2], 'col2': [3,4]}), columns_to_be_checked = ('col1'))
+    >>> check_null(dataframe = pd.DataFrame({'col1': [1,2], 'col2': [3,4]}),columns_to_be_checked = ['col1', 'col2'])
     0
-    >>> check_null(dataframe = pd.DataFrame({'col1': [1,numpy.nan], 'col2': [3,4]}), columns_to_be_checked = ('col1'))
+    >>> check_null(dataframe = pd.DataFrame({'col1': [1,numpy.nan], 'col2': [3,4]}),columns_to_be_checked = ['col1'])
     1
     """
 
-    for eachVal in columns_to_be_checked:
-        if type(eachVal) != str:
-            raise ValueError("Please input string values for column names.")
-    if columns_to_be_checked not in dataframe.columns:
-        raise KeyError("Please check the column names correspond to values in the DataFrame.")
+    if not isinstance(columns_to_be_checked, list):
+        raise ValueError("Please make sure that all your columns passed are strings")
+    
+    for eachCol in columns_to_be_checked:
+        if eachCol not in dataframe.columns:
+            raise KeyError("Please check the column names correspond to values in the DataFrame.")
 
     null_count = 0
     for eachColumn in columns_to_be_checked:
         prev_null_count = null_count
         null_count = prev_null_count + (len(dataframe) - dataframe[eachColumn].count())
-        
+
     return null_count
diff --git a/codonPython/tests/check_consistent_measures_test.py b/codonPython/tests/check_consistent_measures_test.py
@@ -1,5 +1,6 @@
 from codonPython.check_consistent_measures import check_consistent_measures
 import pandas as pd
+import numpy as np
 import pytest
 
 @pytest.mark.parametrize("data, geography_col, measure_col, measures_set, expected", [
@@ -30,34 +31,28 @@
 def test_each_org_levels_BAU(data, geography_col, measure_col, measures_set, expected):
     assert expected == check_consistent_measures(data, geography_col, measure_col, measures_set)
 
-@pytest.mark.parametrize("data, geography_col, measure_col, measures_set", [ 
+
+@pytest.mark.parametrize("data, geography_col, measure_col, measures_set", [
     (
         pd.DataFrame({
-            "Geog" : ["National" ,"National", 1, "Region", "Local", "Local"],
-            "measure" : ["m1", "m2", "m1", "m2", "m1", "m2"],
+            "Geog" : ["National" ,"National", "Region", "Region", "Local", "Local"],
+            "measure" : ["m1", "m2", "m1", np.nan, "m1", "m2"],
             "Value_Unsuppressed" : [4, 2, 2, 1, 2, 1],
         }),
         "Geog",
         "measure",
         set({"m1", "m2"}),
-    )   
-])
-
-def test_each_org_levels_valueErrors(data, geography_col, measure_col, measures_set):
-    with pytest.raises(ValueError):
-        check_consistent_measures(data, geography_col, measure_col, measures_set)
-
-@pytest.mark.parametrize("data, geography_col, measure_col, measures_set", [
+    ),
     (
         pd.DataFrame({
-            "Geog" : ["National" ,"National", "Region", "Region", "Local", "Local"],
-            "measure" : ["m1", "m2", "m1", True, "m1", "m2"],
+            "Geog" : ["National" ,"National", 1, "Region", "Local", "Local"],
+            "measure" : ["m1", "m2", "m1", "m2", "m1", "m2"],
             "Value_Unsuppressed" : [4, 2, 2, 1, 2, 1],
         }),
         "Geog",
         "measure",
         set({"m1", "m2"}),
-    )
+    )   
 ])
 
 
