Python unit test: Database query result verification

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

Write a Python unit test program to check if a database query returns the expected results.

Sample Solution:

Python Code:

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

# Import the 'sqlite3' module for working with SQLite databases.
import sqlite3

# Define a test case class 'TestDatabaseQuery' that inherits from 'unittest.TestCase'.
class TestDatabaseQuery(unittest.TestCase):
    # Define a method 'setUp' that is executed before each test.
    def setUp(self):
        # Create a database connection in memory and insert test data.
        self.conn = sqlite3.connect(':memory:')
        self.cursor = self.conn.cursor()
        # Create an 'employees' table and insert test records.
        self.cursor.execute("CREATE TABLE employees (id INTEGER PRIMARY KEY, name TEXT, salary REAL)")
        self.cursor.execute("INSERT INTO employees (name, salary) VALUES ('Ylva Guiomar', 1800.0)")
        self.cursor.execute("INSERT INTO employees (name, salary) VALUES ('Scott Gregorius', 2100.0)")
        self.conn.commit()

    # Define a method 'tearDown' that is executed after each test.
    def tearDown(self):
        # Close the database cursor and the database connection.
        self.cursor.close()
        self.conn.close()

    # Define a test method 'test_database_query' to test a database query.
    def test_database_query(self):
        # Execute a SQL query to select employee names and salaries, ordered by name.
        self.cursor.execute("SELECT name, salary FROM employees ORDER BY name")
        results = self.cursor.fetchall()

        # Define the expected results as a list of tuples.
        expected_results = [('Scott Gregorius', 2100.0), ('Ylva Guiomar', 1800.0)]

        # Assert that the results match the expected results.
        self.assertEqual(results, expected_results)

# 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 1 test in 0.001s
OK

Explanation:

The unittest module in Python can be used to test if a database query returns the expected results, along with a suitable database library, such as sqlite3, psycopg2, or mysql-connector-python.

In the above exercise,

• The setUp method is called before each test case. Inside this method, we create a database connection to an in-memory SQLite database using sqlite3.connect(':memory:'). We then create a cursor to execute SQL queries on the database.
• In the setUp method, we also create a table called employees and insert some test data. The table has columns id, name, and salary, and we insert two rows into the table.
• After each test case, the tearDown method closes the cursor and the database connection.
• The test_database_query method is the actual test case. Using cursor.execute(), we execute the query "SELECT name, salary FROM employees ORDER BY name" and fetch all the results using cursor.fetchall ().
• We define the expected results as a list of tuples in the expected_results variable. In this example, we expect the names to be ordered alphabetically.
• Finally, we use self.assertEqual(results, expected_results) to assert that the database results match the expected results.
• The unittest module executes the test_database_query method when running the unit test program. As long as the database query returns the expected results, the assertion will pass.

Flowchart:

Previous: Database connection validation.
Next: Input data parsing and validation.