from itertools import count
from unittest import TestCase

from happy_numbers import *
from happy_numbers_recursion import *


class TestHappyNumbers(TestCase):
    def test_next_in_chain(self):
        self.assertEqual(next_in_chain(1), 1)
        self.assertEqual(next_in_chain(19), 82)
        self.assertEqual(next_in_chain(82), 68)

    def test_first_five_happy_numbers(self):
        happy_numbers = (n for n in count(1) if is_happy(n))
        first_five = [next(happy_numbers) for _ in range(5)]
        self.assertEqual(first_five, [1, 7, 10, 13, 19])

    def test_number_of_happy_numbers_under_1001(self):
        happy_numbers = [i for i in range(1001) if is_happy(i)]
        self.assertEqual(len(happy_numbers), 143)


class TestRecursiveHappyNumbers(TestCase):
    def test_number_of_happy_numbers_under_1001(self):
        happy_numbers = [i for i in range(1001) if is_happy_number(i)]
        self.assertEqual(len(happy_numbers), 143)