Allow drawing of multiple cards at once

game/models.py

@@ -87,8 +87,18 @@ class Deck:
     def __init__(self, cards):
         self.cards = cards
 
-    def draw(self):
-        return self.cards.pop()
+    def draw(self, count):
+        """
+        Draw <count> cards from the deck, or as many as are available.
+        return a list of cards drawn.
+        """
+        drawn = []
+        while len(drawn) < count:
+            try:
+                drawn.append(self.cards.pop())
+            except IndexError:
+                break
+        return drawn
 
     def shuffle(self):
         random.shuffle(self.cards)

tests/test_deck.py

@@ -5,8 +5,15 @@ def test_deck_draw_removes_items():
     d = Deck(list(range(16)))
     assert len(d.cards) == 16
     d.shuffle()
-    d.draw()
-    assert len(d.cards) == 15
+    d.draw(4)
+    assert len(d.cards) == 12
+
+
+def test_drawing_past_end_of_deck():
+    d = Deck(list(range(3)))
+    drawn = d.draw(4)
+    assert len(d.cards) == 0
+    assert len(drawn) == 3
 
 
 def test_deck_add():