bpo-29984: Improve 'heapq' test coverage (GH-992)

Author: rkday

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

@@ -2,6 +2,7 @@
 
 import random
 import unittest
+import doctest
 
 from test import support
 from unittest import TestCase, skipUnless
@@ -26,6 +27,23 @@ def test_c_functions(self):
             self.assertEqual(getattr(c_heapq, fname).__module__, '_heapq')
 
 
+def load_tests(loader, tests, ignore):
+    # The 'merge' function has examples in its docstring which we should test
+    # with 'doctest'.
+    #
+    # However, doctest can't easily find all docstrings in the module (loading
+    # it through import_fresh_module seems to confuse it), so we specifically
+    # create a finder which returns the doctests from the merge method.
+
+    class HeapqMergeDocTestFinder:
+        def find(self, *args, **kwargs):
+            dtf = doctest.DocTestFinder()
+            return dtf.find(py_heapq.merge)
+
+    tests.addTests(doctest.DocTestSuite(py_heapq,
+                                        test_finder=HeapqMergeDocTestFinder()))
+    return tests
+
 class TestHeap:
 
     def test_push_pop(self):
@@ -135,6 +153,13 @@ def test_heappushpop(self):
         x = self.module.heappushpop(h, 11)
         self.assertEqual((h, x), ([11], 10))
 
+    def test_heappop_max(self):
+        # _heapop_max has an optimization for one-item lists which isn't
+        # covered in other tests, so test that case explicitly here
+        h = [3, 2]
+        self.assertEqual(self.module._heappop_max(h), 3)
+        self.assertEqual(self.module._heappop_max(h), 2)
+
     def test_heapsort(self):
         # Exercise everything with repeated heapsort checks
         for trial in range(100):
@@ -168,6 +193,12 @@ def test_merge(self):
                                  list(self.module.merge(*seqs, key=key, reverse=reverse)))
                 self.assertEqual(list(self.module.merge()), [])
 
+    def test_empty_merges(self):
+        # Merging two empty lists (with or without a key) should produce
+        # another empty list.
+        self.assertEqual(list(self.module.merge([], [])), [])
+        self.assertEqual(list(self.module.merge([], [], key=lambda: 6)), [])
+
     def test_merge_does_not_suppress_index_error(self):
         # Issue 19018: Heapq.merge suppresses IndexError from user generator
         def iterable():