Udpate test_heapq from CPtyhon 3.10.6

Author: CPython Developers

Lib/heapq.py

@@ -597,5 +597,5 @@ def nlargest(n, iterable, key=None):
 
 if __name__ == "__main__":
 
-    import doctest
-    print(doctest.testmod())
+    import doctest # pragma: no cover
+    print(doctest.testmod()) # pragma: no cover

Lib/test/test_heapq.py

@@ -146,11 +146,11 @@ def test_heappushpop(self):
         self.assertEqual(type(h[0]), int)
         self.assertEqual(type(x), float)
 
-        h = [10];
+        h = [10]
         x = self.module.heappushpop(h, 9)
         self.assertEqual((h, x), ([10], 9))
 
-        h = [10];
+        h = [10]
         x = self.module.heappushpop(h, 11)
         self.assertEqual((h, x), ([11], 10))
 
@@ -433,6 +433,37 @@ def test_heappop_mutating_heap(self):
         with self.assertRaises((IndexError, RuntimeError)):
             self.module.heappop(heap)
 
+    def test_comparison_operator_modifiying_heap(self):
+        # See bpo-39421: Strong references need to be taken
+        # when comparing objects as they can alter the heap
+        class EvilClass(int):
+            def __lt__(self, o):
+                heap.clear()
+                return NotImplemented
+
+        heap = []
+        self.module.heappush(heap, EvilClass(0))
+        self.assertRaises(IndexError, self.module.heappushpop, heap, 1)
+
+    def test_comparison_operator_modifiying_heap_two_heaps(self):
+
+        class h(int):
+            def __lt__(self, o):
+                list2.clear()
+                return NotImplemented
+
+        class g(int):
+            def __lt__(self, o):
+                list1.clear()
+                return NotImplemented
+
+        list1, list2 = [], []
+
+        self.module.heappush(list1, h(0))
+        self.module.heappush(list2, g(0))
+
+        self.assertRaises((IndexError, RuntimeError), self.module.heappush, list1, g(1))
+        self.assertRaises((IndexError, RuntimeError), self.module.heappush, list2, h(1))
 
 class TestErrorHandlingPython(TestErrorHandling, TestCase):
     module = py_heapq