15. Armies Battle / Вирішуємо задачі на Python CheckIO Українською

# -------------   tasks/wariors.py   ----------------------


class Warrior:
    health = 50
    attack = 5
    defense = 0

    @property
    def is_alive(self) -> bool:
        return self.health > 0

    def deal_dmg(self, unit: "Warrior"):
        unit.health -= max(0, self.attack - unit.defense)


class Knight(Warrior):
    attack = 7


class Defender(Warrior):
    health = 60
    attack = 3
    defense = 2


def fight(unit_1: Warrior, unit_2: Warrior) -> bool:

    while unit_1.is_alive:
        unit_1.deal_dmg(unit_2)
        if not unit_2.is_alive:
            break
        unit_2.deal_dmg(unit_1)

    return unit_1.is_alive


# -----------------     tasks/army.py     -------------------------


from collections import UserList

from tasks.wariors import fight


class Army(UserList):

    def add_units(self, unit_type: type, count: int):
        self.data.extend(unit_type() for _ in range(count))


class Battle:

    def fight(self, army1: Army, army2: Army) -> bool:

        while army1 and army2:
            unit1, unit2 = army1[0], army2[0]
            if fight(unit1, unit2):
                army2.pop(0)
            else:
                army1.pop(0)
        return bool(army1)



# -----------     tests/test_wariors.py     ----------------------

from tasks.wariors import Warrior, Knight, fight, Defender
from tasks.army import Army, Battle


def test_all():

    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

    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

    # battle tests
    my_army = Army()
    my_army.add_units(Knight, 3)

    enemy_army = Army()
    enemy_army.add_units(Warrior, 3)

    army_3 = Army()
    army_3.add_units(Warrior, 20)
    army_3.add_units(Knight, 5)

    army_4 = Army()
    army_4.add_units(Warrior, 30)

    battle = Battle()

    assert battle.fight(my_army, enemy_army) == True
    assert battle.fight(army_3, army_4) == False

    chuck = Warrior()
    bruce = Warrior()
    carl = Knight()
    dave = Warrior()
    mark = Warrior()
    bob = Defender()
    mike = Knight()
    rog = Warrior()
    lancelot = Defender()

    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
    assert fight(bob, mike) == False
    assert fight(lancelot, rog) == True

    my_army = Army()
    my_army.add_units(Defender, 1)

    enemy_army = Army()
    enemy_army.add_units(Warrior, 2)

    army_3 = Army()
    army_3.add_units(Warrior, 1)
    army_3.add_units(Defender, 1)

    army_4 = Army()
    army_4.add_units(Warrior, 2)

    battle = Battle()

    assert battle.fight(my_army, enemy_army) == False
    assert battle.fight(army_3, army_4) == True