Why the fourth test fails. With army1 and with army2, for some reason, Army 1 wins, although army 2 has a whole unit in reserve.

I came to the conclusion that the unit of the first army, after killing an enemy unit, beats the next enemy unit, and the unit of the second army does not have such an opportunity.

At the end of the task clarification - the first army has the advantage to start each battle. So I understand it means every round of the battle and not the battle itself. Right?

Code:

class Army:
    def __init__(self):
        self.units = []

    def add_units(self, unit, count):
        self.units.extend([unit for i in range(count)])


class Warrior:
    def __init__(self):
        self.health_points = 50
        self.attack_points = 5

    def take_damage(self, damage):
        self.health_points -= damage

    def an_attack(self, enemy_health_point):
        enemy_health_point.take_damage(self.attack_points)

    @property
    def is_alive(self):
        return self.health_points > 0


class Knight(Warrior):
    def __init__(self):
        super(Knight, self).__init__()
        self.attack_points = 7


class Battle:

    def fight(self, army_1, army_2):
        self.army_1 = iter(army_1.units)
        self.army_2 = iter(army_2.units)

        def next_unit(army):
            return next(army)()

        unit_1 = next_unit(self.army_1)
        unit_2 = next_unit(self.army_2)
        count = 0
        while True:
            if unit_1.is_alive:
                unit_1.an_attack(unit_2)
            else:
                try:
                    unit_1 = next_unit(self.army_1)
                except StopIteration:
                    return False

            if unit_2.is_alive:
                unit_2.an_attack(unit_1)
            else:
                try:
                    unit_2 = next_unit(self.army_2)
                except StopIteration:
                    return True


def fight(unit_1, unit_2):
    while True:
        unit_1.an_attack(unit_2)
        if not unit_2.is_alive:
            return True
        unit_2.an_attack(unit_1)
        if not unit_1.is_alive:
            return False


if __name__ == '__main__':
    chuck = Warrior()
    bruce = Warrior()
    carl = Knight()
    dave = Warrior()
    mark = Warrior()

    assert fight(chuck, bruce) == True
    assert fight(dave, carl) == False
    assert chuck.is_alive == True
    assert bruce.is_alive == False
    assert carl.is_alive == True
    assert dave.is_alive == False
    assert fight(carl, mark) == False
    assert carl.is_alive == False

    army_1 = Army()
    army_2 = Army()
    army_1.add_units(Warrior, 20)
    army_2.add_units(Warrior, 21)
    battle = Battle()
    assert battle.fight(army_1, army_2) == True