import unittest
try:
    from .lib_prime import prime_factors
    from .lib_prime import prime_factors_count
    from .lib_prime import is_prime
    from .lib_prime import prime_nth
    from .lib_prime import primes
    from .lib_prime import get_divisors_count
except ModuleNotFoundError:
    from lib_prime import prime_factors
    from lib_prime import prime_factors_count
    from lib_prime import is_prime
    from lib_prime import prime_nth
    from lib_prime import primes
    from lib_prime import get_divisors_count


class TestPrimeMethods(unittest.TestCase):

    def test_prime_factors(self):
        self.assertEqual(prime_factors(2), [2])
        self.assertEqual(prime_factors(5), [5])
        self.assertEqual(prime_factors(10), [2, 5])
        self.assertEqual(prime_factors(13), [13])
        self.assertEqual(prime_factors(147), [3, 7, 7])

    def test_prime_factors_count(self):
        self.assertEqual(prime_factors_count(1), {})
        self.assertEqual(prime_factors_count(2), {2: 1})
        self.assertEqual(prime_factors_count(147), {3: 1, 7: 2})

    def test_is_prime(self):
        self.assertTrue(is_prime(2))
        self.assertTrue(is_prime(3))
        self.assertTrue(is_prime(5))
        self.assertTrue(is_prime(7))
        self.assertTrue(is_prime(11))
        self.assertTrue(is_prime(13))
        self.assertTrue(is_prime(101))

        self.assertFalse(is_prime(0))
        self.assertFalse(is_prime(1))
        self.assertFalse(is_prime(4))
        self.assertFalse(is_prime(6))
        self.assertFalse(is_prime(8))
        self.assertFalse(is_prime(9))
        self.assertFalse(is_prime(10))
        self.assertFalse(is_prime(12))
        self.assertFalse(is_prime(323))

    def test_get_nth_prime(self):
        self.assertEqual(prime_nth(1), 2)
        self.assertEqual(prime_nth(2), 3)
        self.assertEqual(prime_nth(3), 5)
        self.assertEqual(prime_nth(4), 7)
        self.assertEqual(prime_nth(5), 11)
        self.assertEqual(prime_nth(6), 13)

    def test_primes(self):
        self.assertEqual(primes(19), [2, 3, 5, 7, 11, 13, 17, 19])
        self.assertEqual(primes(20), [2, 3, 5, 7, 11, 13, 17, 19])
        self.assertEqual(primes(25), [2, 3, 5, 7, 11, 13, 17, 19, 23])
        self.assertEqual(primes(1), [])

    def test_get_divisors_count(self):
        self.assertEqual(get_divisors_count(1), 1)
        self.assertEqual(get_divisors_count(3), 2)
        self.assertEqual(get_divisors_count(6), 4)
        self.assertEqual(get_divisors_count(10), 4)
        self.assertEqual(get_divisors_count(15), 4)
        self.assertEqual(get_divisors_count(21), 4)
        self.assertEqual(get_divisors_count(28), 6)


if __name__ == '__main__':
    unittest.main()