Python unit test: Check prime number

Last update on December 21 2024 07:48:48 (UTC/GMT +8 hours)

Write a Python unit test program to check if a given number is prime or not.

Sample Solution:

Python Code:

# Import the 'unittest' module for writing unit tests.
import unittest

# Define a function 'is_prime' to check if a number is prime.
def is_prime(number):
    if number < 2:
        return False
    for i in range(2, int(number**0.5) + 1):
        if number % i == 0:
            return False
    return True

# Define a test case class 'PrimeNumberTestCase' that inherits from 'unittest.TestCase'.
class PrimeNumberTestCase(unittest.TestCase):
    # Define a test method 'test_prime_numbers' to test prime numbers.
    def test_prime_numbers(self):
        prime_numbers = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]
        # You can use the alternative 'prime_numbers' list by uncommenting it.
        # prime_numbers = [2, 3, 4, 8, 11, 13, 17, 19, 23, 30, 31]
        print("Prime numbers:", prime_numbers)
        # Iterate through prime numbers and assert that they are recognized as prime.
        for number in prime_numbers:
            self.assertTrue(is_prime(number), f"{number} is not recognized as a prime number")

    # Define a test method 'test_non_prime_numbers' to test non-prime numbers.
    def test_non_prime_numbers(self):
        non_prime_numbers = [4, 6, 8, 10, 12, 14, 16, 18, 20]
        # You can use the alternative 'non_prime_numbers' list by uncommenting it.
        # non_prime_numbers = [4, 6, 8, 9, 11, 12, 14, 17, 16, 18, 20]
        print("Non-prime numbers:", non_prime_numbers)
        # Iterate through non-prime numbers and assert that they are not recognized as prime.
        for number in non_prime_numbers:
            self.assertFalse(is_prime(number), f"{number} is incorrectly recognized as a prime number")

# Check if the script is run as the main program.
if __name__ == '__main__':
    # Run the test cases using 'unittest.main()'.
    unittest.main()

Sample Output:

Ran 2 tests in 0.002s
OK
Non prime numbers: [4, 6, 8, 10, 12, 14, 16, 18, 20]
Prime numbers: [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31]

Ran 2 tests in 0.002s
FAILED (failures=2)
Non prime numbers: [4, 6, 8, 9, 11, 12, 14, 17, 16, 18, 20]
Prime numbers: [2, 3, 4, 8, 11, 13, 17, 19, 23, 30, 31]

Explanation:

In the above exercise,

We define a function is_prime() that checks if a given number is prime. Then, we create a test case class PrimeNumberTestCase that subclasses unittest.TestCase. Inside this class, we define multiple test methods using the test_ prefix. Each test method uses the assertion methods provided by unittest.TestCase to check if the expected condition is true or false.

Flowchart:

Previous: Python Unit test Exercises Home.
Next: Check list sorting.