Merge pull request #171 from datarian/fix-userwarning-unittests

Fix userwarning unittests

Author: janmotl

category_encoders/tests/test_basen.py

@@ -111,14 +111,13 @@ def test_inverse_transform_BothFieldsAreReturnNanWithNan_ExpectValueError(self):
         enc = encoders.BaseNEncoder(handle_missing='return_nan', handle_unknown='return_nan')
         enc.fit(train)
         result = enc.transform(test)
+        
+        message = 'inverse_transform is not supported because transform impute '\
+                  'the unknown category nan when encode city'
 
-        with warnings.catch_warnings(record=True) as w:
+        with self.assertWarns(UserWarning, msg=message) as w:
             enc.inverse_transform(result)
 
-            self.assertEqual(1, len(w))
-            self.assertEqual('inverse_transform is not supported because transform impute '
-                             'the unknown category nan when encode city', str(w[0].message))
-
     def test_inverse_transform_HaveMissingAndNoUnknown_ExpectInversed(self):
         train = pd.DataFrame({'city': ['chicago', np.nan]})
         test = pd.DataFrame({'city': ['chicago', 'los angeles']})

category_encoders/tests/test_one_hot.py

@@ -227,14 +227,13 @@ def test_inverse_transform_BothFieldsAreReturnNanWithNan_ExpectValueError(self):
         enc = encoders.OneHotEncoder(handle_missing='return_nan', handle_unknown='return_nan')
         enc.fit(train)
         result = enc.transform(test)
+        
+        message = 'inverse_transform is not supported because transform impute '\
+                  'the unknown category nan when encode city'
 
-        with warnings.catch_warnings(record=True) as w:
+        with self.assertWarns(UserWarning, msg=message) as w:
             enc.inverse_transform(result)
 
-            self.assertEqual(1, len(w))
-            self.assertEqual('inverse_transform is not supported because transform impute '
-                             'the unknown category nan when encode city', str(w[0].message))
-
     def test_inverse_transform_HaveMissingAndNoUnknown_ExpectInversed(self):
         train = pd.DataFrame({'city': ['chicago', np.nan]})
         test = pd.DataFrame({'city': ['chicago', 'los angeles']})

category_encoders/tests/test_ordinal.py

@@ -137,13 +137,12 @@ def test_inverse_transform_HaveUnknown_ExpectWarning(self):
         enc.fit(train)
         result = enc.transform(test)
 
-        with warnings.catch_warnings(record=True) as w:
+        message = 'inverse_transform is not supported because transform impute '\
+                             'the unknown category -1 when encode city'
+        
+        with self.assertWarns(UserWarning, msg=message) as w:
             enc.inverse_transform(result)
 
-            self.assertEqual(1, len(w))
-            self.assertEqual('inverse_transform is not supported because transform impute '
-                             'the unknown category -1 when encode city', str(w[0].message))
-
     def test_inverse_transform_HaveNanInTrainAndHandleMissingValue_ExpectReturnedWithNan(self):
         train = pd.DataFrame({'city': ['chicago', np.nan]})
 
@@ -169,14 +168,13 @@ def test_inverse_transform_BothFieldsAreReturnNanWithNan_ExpectValueError(self):
         enc = encoders.OrdinalEncoder(handle_missing='return_nan', handle_unknown='return_nan')
         enc.fit(train)
         result = enc.transform(test)
+        
+        message = 'inverse_transform is not supported because transform impute '\
+                             'the unknown category nan when encode city'
 
-        with warnings.catch_warnings(record=True) as w:
+        with self.assertWarns(UserWarning, msg=message) as w:
             enc.inverse_transform(result)
 
-            self.assertEqual(1, len(w))
-            self.assertEqual('inverse_transform is not supported because transform impute '
-                             'the unknown category nan when encode city', str(w[0].message))
-
     def test_inverse_transform_HaveMissingAndNoUnknown_ExpectInversed(self):
         train = pd.DataFrame({'city': ['chicago', np.nan]})
         test = pd.DataFrame({'city': ['chicago', 'los angeles']})