primes.py

import unittest


def is_prime(n):
    if n <= 1:
        return False

    if n == 2 or n == 3:
        return True

    if n % 2 == 0 or n % 3 == 0:
        return False

    divisor = 6
    while n >= divisor*divisor - 2*divisor + 1:
        if n % (divisor - 1) == 0:
            return False

        if n % (divisor + 1) == 0:
            return False

        divisor += 6

    return True


def find_next_prime(n):
    n += 1
    while not is_prime(n):
        n += 1
    return n


class TestNextPrime(unittest.TestCase):
    def test_next_prime(self):
        self.assertEqual(find_next_prime(2), 3)
        self.assertEqual(find_next_prime(3), 5)
        self.assertEqual(find_next_prime(6), 7)
        self.assertEqual(find_next_prime(10), 11)
        self.assertEqual(find_next_prime(13), 17)
        self.assertEqual(find_next_prime(788), 797)
        self.assertEqual(find_next_prime(1), 2)
        self.assertEqual(find_next_prime(-10), 2)
        self.assertEqual(find_next_prime(-123), 2)
        self.assertEqual(find_next_prime(321312), 321313)
        self.assertEqual(find_next_prime(370247), 370261)


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