SF patch #969791: Add nlargest() and nsmallest() to heapq.

1 files changed, 10 insertions, 1 deletions

diff --git a/Lib/test/test_heapq.py b/Lib/test/test_heapq.py
index 8f3c6f9..f04ea94 100644
--- a/Lib/test/test_heapq.py
+++ b/Lib/test/test_heapq.py
@@ -2,7 +2,7 @@
 
 from test.test_support import verify, vereq, verbose, TestFailed
 
-from heapq import heappush, heappop, heapify, heapreplace
+from heapq import heappush, heappop, heapify, heapreplace, nlargest, nsmallest
 import random
 
 def check_invariant(heap):
@@ -84,6 +84,15 @@ def test_main():
         data.sort()
         sorted = [heappop(heap) for i in range(size)]
         vereq(data, sorted)
+
+    # 7) Check nlargest() and nsmallest()
+    data = [random.randrange(2000) for i in range(1000)]
+    copy = data[:]
+    copy.sort(reverse=True)
+    vereq(nlargest(data, 400), copy[:400])
+    copy.sort()
+    vereq(nsmallest(data, 400), copy[:400])
+
     # Make user happy
     if verbose:
         print "All OK"