Use ChatGPT to Write Unit Tests

Contributed by: Mike Loukides

Difficulty level: Intermediate

“Use ChatGPT to Improve Code” gave us a program that was initially written by a human, but modified by ChatGPT. Can we get ChatGPT to write unit tests for the program? Yes we can:

Prompt: Write a unit test for the function preprocess_data.

ChatGPT: To write a unit test for the preprocess_data function, you can use the unittest module in Python. First, import the necessary modules and the preprocess_data function from your main script. Then, create a class that inherits from unittest.TestCase and define test methods for various cases. Here’s an example of a unit test for the preprocess_data function:

import unittest
import pandas as pd
from your_main_script import preprocess_data

class TestPreprocessData(unittest.TestCase):

    def test_preprocess_data(self):
        # Create a sample DataFrame with missing values
        raw_data = {
            'A': [1, 4, 7],
            'B': [None, 5, 8],
            'C': [3, 6, None]
        }
        df = pd.DataFrame(raw_data)

        # Call the preprocess_data function
        preprocess_data(df)

        # Expected DataFrame after preprocessing
        expected_data = {
            'A': [1, 4, 7],
            'B': ['No answer', 5, 8],
            'C': [3, 6, 'No answer']
        }
        expected_df = pd.DataFrame(expected_data)

        # Compare the preprocessed DataFrame with the
        # expected one
        pd.testing.assert_frame_equal(df, expected_df)

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

Here’s something to watch for, though. I lied. This isn’t the first unit test that ChatGPT wrote. The first one used Python’s ...