Add regression test for add() after remove() with hash collision in set (GH-143781)

diff --git a/Lib/test/test_dict.py b/Lib/test/test_dict.py
index 3b4e95015e5b4c594f63b9fc55a4b8fe56daa642..6583c0f2aefb2b670015facea702c076b89bcfac 100644 (file)
--- a/Lib/test/test_dict.py
+++ b/Lib/test/test_dict.py
@@ -12,6 +12,15 @@ from test import support
 from test.support import import_helper
 
 
+class CustomHash:
+    def __init__(self, hash):
+        self.hash = hash
+    def __hash__(self):
+        return self.hash
+    def __repr__(self):
+        return f'<CustomHash {self.hash} at {id(self):#x}>'
+
+
 class DictTest(unittest.TestCase):
 
     def test_invalid_keyword_arguments(self):
@@ -1648,6 +1657,29 @@ class DictTest(unittest.TestCase):
         d[MyStr("attr1")] = 2
         self.assertIsInstance(list(d)[0], MyStr)
 
+    def test_hash_collision_remove_add(self):
+        self.maxDiff = None
+        # There should be enough space, so all elements with unique hash
+        # will be placed in corresponding cells without collision.
+        n = 64
+        items = [(CustomHash(h), h) for h in range(n)]
+        # Keys with hash collision.
+        a = CustomHash(n)
+        b = CustomHash(n)
+        items += [(a, 'a'), (b, 'b')]
+        d = dict(items)
+        self.assertEqual(len(d), len(items), d)
+        del d[a]
+        # "a" has been replaced with a dummy.
+        del items[n]
+        self.assertEqual(len(d), len(items), d)
+        self.assertEqual(d, dict(items))
+        d[b] = 'c'
+        # "b" should not replace the dummy.
+        items[n] = (b, 'c')
+        self.assertEqual(len(d), len(items), d)
+        self.assertEqual(d, dict(items))
+
 
 class CAPITest(unittest.TestCase):
 

diff --git a/Lib/test/test_set.py b/Lib/test/test_set.py
index 3539adc789c7e712d86b6487bb1064de862feff2..554716aed1e98b7cd8a96277e35036f439acce66 100644 (file)
--- a/Lib/test/test_set.py
+++ b/Lib/test/test_set.py
@@ -20,6 +20,14 @@ def check_pass_thru():
     raise PassThru
     yield 1
 
+class CustomHash:
+    def __init__(self, hash):
+        self.hash = hash
+    def __hash__(self):
+        return self.hash
+    def __repr__(self):
+        return f'<CustomHash {self.hash} at {id(self):#x}>'
+
 class BadCmp:
     def __hash__(self):
         return 1
@@ -675,6 +683,28 @@ class TestSet(TestJointOps, unittest.TestCase):
         with self.assertRaises(KeyError):
             myset.discard(elem2)
 
+    def test_hash_collision_remove_add(self):
+        self.maxDiff = None
+        # There should be enough space, so all elements with unique hash
+        # will be placed in corresponding cells without collision.
+        n = 64
+        elems = [CustomHash(h) for h in range(n)]
+        # Elements with hash collision.
+        a = CustomHash(n)
+        b = CustomHash(n)
+        elems += [a, b]
+        s = self.thetype(elems)
+        self.assertEqual(len(s), len(elems), s)
+        s.remove(a)
+        # "a" has been replaced with a dummy.
+        del elems[n]
+        self.assertEqual(len(s), len(elems), s)
+        self.assertEqual(s, set(elems))
+        s.add(b)
+        # "b" should not replace the dummy.
+        self.assertEqual(len(s), len(elems), s)
+        self.assertEqual(s, set(elems))
+
 
 class SetSubclass(set):
     pass