from unittest import TestCase
from benchmark.check_armstrong import check_armstrong


class Test_check_armstrong(TestCase):
    #  distances_true = {1: [0]}
    #  distances_false = {1: [1]}
    def test_check_armstrong_1(self):
        assert check_armstrong(n=0) == True

    #  distances_true = {1: [2], 2: [1], 3: [0]}
    #  distances_false = {1: [0], 2: [0], 3: [149]}
    def test_check_armstrong_2(self):
        assert check_armstrong(n=2) == False

    #  distances_true = {1: [1], 2: [0]}
    #  distances_false = {1: [0], 2: [1]}
    def test_check_armstrong_3(self):
        assert check_armstrong(n=1) == True

    #  distances_true = {1: [150], 2: [149], 3: [0]}
    #  distances_false = {1: [0], 2: [0], 3: [1]}
    def test_check_armstrong_4(self):
        assert check_armstrong(n=150) == False

    #  distances_true = {1: [151], 2: [150], 3: [1], 4: [0], 5: [151]}
    #  distances_false = {1: [0], 2: [0], 3: [0], 4: [151], 5: [0]}
    def test_check_armstrong_5(self):
        assert check_armstrong(n=151) == False