diff --git a/logic/__tests__/knights_queue_test.py b/logic/__tests__/knights_queue_test.py index 99ba7da..ecf9548 100644 --- a/logic/__tests__/knights_queue_test.py +++ b/logic/__tests__/knights_queue_test.py @@ -5,6 +5,52 @@ from models.knight import Knight class KnightsQueueTest(unittest.TestCase): + def test_should_skip_dead_knights(self): + knight1 = Knight(None) + knight1.health = 0 + + knight2 = Knight(None) + knight2.health = 0 + + knight3 = Knight(None) + knight3.health = 1 + + knight4 = Knight(None) + knight4.health = 0 + + knight5 = Knight(None) + knight5.health = 0 + + knight6 = Knight(None) + knight6.health = 1 + + knights_queue = KnightsQueue([knight1, knight2, knight3], [knight4, knight5, knight6]) + + res1 = knights_queue.dequeue_knight() + res2 = knights_queue.dequeue_knight() + + self.assertEqual(res1.health, 1) + self.assertEqual(res2.health, 1) + + def test_should_return_first_alive_knight(self): + knight1 = Knight(None) + knight1.health = 222 + + knight2 = Knight(None) + knight2.health = -1 + + knight3 = Knight(None) + knight3.health = 1 + + knights_queue = KnightsQueue([knight1, knight2], [knight3]) + + res1 = knights_queue.dequeue_knight() + res2 = knights_queue.dequeue_knight() + res3 = knights_queue.dequeue_knight() + res4 = knights_queue.dequeue_knight() + + self.assertEqual(res1, res3) + self.assertEqual(res2, res4) def test_should_raise_when_knight_died_and_whole_team_dead(self): with self.assertRaises(Exception): @@ -18,6 +64,7 @@ class KnightsQueueTest(unittest.TestCase): knights_queue.dequeue_knight() knight2.health = -2 knights_queue.dequeue_knight() + knights_queue.dequeue_knight() def test_should_make_valid_next_turn(self): knight1 = Knight(None) diff --git a/logic/knights_queue.py b/logic/knights_queue.py index cb3ba2c..4b63e1f 100644 --- a/logic/knights_queue.py +++ b/logic/knights_queue.py @@ -11,7 +11,7 @@ class KnightsQueue: if self.both_teams_alive(): knight = self.queues[self.team_idx_turn].popleft() if knight.health <= 0: - self.dequeue_knight() + return self.dequeue_knight() else: self.queues[self.team_idx_turn].append(knight) self.next_turn()