From cf028c24cfa8b8b4787aea50ad73cc8b18f15770 Mon Sep 17 00:00:00 2001
From: Leonard Richardson <leonardr@segfault.org>
Date: Mon, 26 Aug 2019 18:36:43 -0400
Subject: It's now possible to override any of the element classes.

---
 bs4/tests/test_soup.py | 23 +++++++++++++++--------
 1 file changed, 15 insertions(+), 8 deletions(-)

(limited to 'bs4/tests/test_soup.py')

diff --git a/bs4/tests/test_soup.py b/bs4/tests/test_soup.py
index 07f29c4..af5f791 100644
--- a/bs4/tests/test_soup.py
+++ b/bs4/tests/test_soup.py
@@ -13,6 +13,7 @@ from bs4 import (
 )
 from bs4.element import (
     CharsetMetaAttributeValue,
+    Comment,
     ContentMetaAttributeValue,
     SoupStrainer,
     NamespacedAttribute,
@@ -123,24 +124,30 @@ class TestConstructor(SoupTest):
             self.assertEqual(" a class ", a['class'])
 
     def test_replacement_classes(self):
-        # Test the ability to pass in replacements for the Tag and
-        # NavigableString class, which will be used when building
-        # the tree.
+        # Test the ability to pass in replacements for element classes
+        # which will be used when building the tree.
         class TagPlus(Tag):
             pass
 
         class StringPlus(NavigableString):
             pass
 
+        class CommentPlus(Comment):
+            pass
+        
         soup = self.soup(
-            "<a><b>foo</b>bar</a>",
-            tag_class=TagPlus, string_class=StringPlus
+            "<a><b>foo</b>bar</a><!--whee-->",
+            element_classes = {
+                Tag: TagPlus,
+                NavigableString: StringPlus,
+                Comment: CommentPlus,
+            }
         )
 
-        # The tree was built with TagPlus and StringPlus objects,
-        # rather than Tag and String objects.
+        # The tree was built with TagPlus, StringPlus, and CommentPlus objects,
+        # rather than Tag, String, and Comment objects.
         assert all(
-            isinstance(x, (TagPlus, StringPlus))
+            isinstance(x, (TagPlus, StringPlus, CommentPlus))
             for x in soup.recursiveChildGenerator()
         )
         
-- 
cgit v1.2.3