Testing Your Code

Introduction

In this lab, we will learn how to write unit tests for our Python code using the built-in unittest module. We will start with a simple function and then build up to testing more complex code, including code that interacts with databases.

Prerequisites

Before starting this lab, you should have a basic understanding of Python programming language. You should also have some familiarity with writing functions in Python.

Skills Graph

Add Numbers

Setup

To get started, we need to create a new Python file for our tests. Let's call it test_my_code.py. We will also create a module that contains the code we want to test. Let's call this module my_code.py.

Simple Test

Let's start with a simple function that returns the sum of two numbers. In my_code.py, create a function named add_numbers that takes two arguments and returns their sum.

def add_numbers(a, b):
    return a + b

Now let's write a test for this function in test_my_code.py. First, we need to import the unittest module and the add_numbersfunction.

import unittest
from my_code import add_numbers

Next, we will create a class named TestAddNumbers that inherits from unittest.TestCase.

class TestAddNumbers(unittest.TestCase):
    pass

Inside this class, we will create a method named test_add_numbers that will test the add_numbers function. We will use the assertEqual method to check that the function returns the correct sum.

class TestAddNumbers(unittest.TestCase):
    def test_add_numbers(self):
        result = add_numbers(2, 3)
        self.assertEqual(result, 5)

Now we can run our tests using the following command:

python -m unittest test_my_code.py

If the test passes, you should see output like the following:

...
----------------------------------------------------------------------
Ran 1 test in 0.000s

OK

The . indicates that one test passed.

Testing Exceptions

Next, let's test a function that raises an exception. In my_code.py, create a function named divide that takes two arguments and returns their quotient. If the second argument is zero, the function should raise a ZeroDivisionError.

def divide(a, b):
    if b == 0:
        raise ZeroDivisionError('Cannot divide by zero')
    return a / b

Now let's write a test for this function in test_my_code.py. We will use the assertRaises method to check that the function raises a ZeroDivisionError when the second argument is zero.

class TestDivide(unittest.TestCase):
    def test_divide_by_zero(self):
        with self.assertRaises(ZeroDivisionError):
            divide(1, 0)

Now we can run our tests again:

python -m unittest test_my_code.py

If the test passes, you should see output like the following:

...
----------------------------------------------------------------------
Ran 2 tests in 0.000s

OK

The .. indicates that two tests passed.

Testing with Databases

Finally, let's test some code that interacts with a database. For this example, we will use SQLite and the sqlite3 module.

In my_code.py, create a function named create_table that creates a new table in an SQLite database.

import sqlite3

def create_table():
    conn = sqlite3.connect('my_database.db')
    c = conn.cursor()
    c.execute('CREATE TABLE IF NOT EXISTS my_table (id INTEGER PRIMARY KEY, name TEXT)')
    conn.commit()
    conn.close()

In this code, we create a new SQLite database table named "my_table" with two columns: "id" and "name". We then define a test class named TestCreateTable with a setUp method that creates an SQLite database for testing.

Now let's write a test for the create_table function. First, we need to create a new database for testing. We will do this in a setup method that runs before each test.

class TestCreateTable(unittest.TestCase):
    def setUp(self):
        self.conn = sqlite3.connect('my_database.db')

    def test_create_table(self):
        create_table()
        c = self.conn.cursor()
        c.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='my_table'")
        result = c.fetchone()
        self.assertIsNotNone(result)

In the setUp method, we create a new SQLite database.

In the test_create_table method, we call the create_table function and then check that the table was created using a SQL query.

Now we can run our tests again:

python -m unittest test_my_code.py

If the test passes, you should see output like the following:

...
----------------------------------------------------------------------
Ran 3 tests in 0.001s

OK

The ... indicates that three tests passed.

Summary

In this lab, we learned how to write unit tests for our Python code using the unittest module. We started with a simple function and then built up to testing more complex code, including code that interacts with databases. With this knowledge, you can write tests for your Python code to ensure that it works as expected and catches bugs early in the development process.