Add test case for whitelist

snguyenthanh · Mar 31, 2020 · 4b2454c · 4b2454c
1 parent a3cc5ff
commit 4b2454c
Showing 1 changed file with 16 additions and 0 deletions.
diff --git a/tests.py b/tests.py
@@ -165,6 +165,22 @@ def test_unicode_vietnamese_2(self):
         self.assertEqual(profanity.censor(bad_text), censored_text)
 
 
+class ProfanityWhitelistTest(unittest.TestCase):
+    def setUp(self):
+        self.maxDiff = None
+        # Pre-load CENSOR_WORDSET
+        profanity.load_censor_words()
+
+    def test_whitelist_words(self):
+        bad_text = "I am gay"
+        censored_text = "I am ****"
+        self.assertEqual(profanity.censor(bad_text), censored_text)
+
+        # Whitelist the word `gay`
+        profanity.load_censor_words(whitelist_words=["gay"])
+        self.assertEqual(profanity.censor(bad_text), bad_text)
+
+
 class ProfanityFileTest(unittest.TestCase):
     def setUp(self):
         self.maxDiff = None