fix tests

appeler · Aug 17, 2023 · aeb5d81 · aeb5d81
1 parent 0534b57
commit aeb5d81
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 61 deletions.
diff --git a/ethnicolr2/tests/test_020_pred_census_ln.py b/ethnicolr2/tests/test_020_pred_census_ln.py
@@ -10,7 +10,7 @@
 import shutil
 import unittest
 import pandas as pd
-from ethnicolr2.pred_census_last_name import pred_census_last_name
+from ethnicolr2 import pred_census_last_name
 from . import capture
 
 race = ["api", "black", "hispanic", "white"]
@@ -29,26 +29,8 @@ def tearDown(self):
         pass
 
     def test_pred_census_ln_2010(self):
-        odf = pred_census_last_name(self.df, "last", 2010)
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
-
-    def test_pred_census_ln_2010_mean(self):
-        odf = pred_census_last_name(self.df, "last", 2010, conf_int=0.9)
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race_mean]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
+        odf = pred_census_last_name(self.df, "last")
+        self.assertTrue(all(odf.true_race == odf.preds))
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/ethnicolr2/tests/test_040_pred_fl.py b/ethnicolr2/tests/test_040_pred_fl.py
@@ -10,8 +10,8 @@
 import shutil
 import unittest
 import pandas as pd
-from ethnicolr2.pred_fl_last_name import pred_fl_last_name
-from ethnicolr2.pred_fl_full_name import pred_fl_full_name
+from ethnicolr2 import pred_fl_last_name
+from ethnicolr2 import pred_fl_full_name
 from pkg_resources import resource_filename
 
 from . import capture
@@ -37,47 +37,11 @@ def tearDown(self):
 
     def test_pred_fl_reg_ln(self):
         odf = pred_fl_full_name(self.df, "last")
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
+        self.assertTrue(all(odf.true_race == odf.preds))
 
     def test_pred_fl_reg_name(self):
         odf = pred_fl_full_name(self.df, "last", "first")
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
-
-    def test_pred_fl_reg_ln_mean(self):
-        odf = pred_fl_last_name(self.df, "last", conf_int=0.9)
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race_mean]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
-
-    def test_pred_fl_reg_name_mean(self):
-        odf = pred_fl_full_name(self.df, "last", "first", conf_int=0.9)
-        self.assertTrue(
-            all(
-                odf[[col for col in odf.columns
-                     if col in race_mean]].sum(axis=1).round(1)
-                == 1.0
-            )
-        )
-        self.assertTrue(all(odf.true_race == odf.race))
+        self.assertTrue(all(odf.true_race == odf.preds))
 
 if __name__ == "__main__":
     unittest.main()