Generate Unit Tests skeleton from Code using OpenAI API

Hey there, Darian Vance here. As a Senior DevOps Engineer at TechResolve, I’m always looking for ways to trim the fat from our development lifecycle. One of the biggest time sinks has always been writing boilerplate code. I used to spend a good chunk of my sprint kickoff just scaffolding out basic unit tests for new services. It’s necessary work, but it’s repetitive. That’s when I started experimenting with the OpenAI API, and I found a workflow that saves me a few hours every week. It doesn’t write perfect tests, but it generates a rock-solid skeleton that gets me 80% of the way there in seconds.

This isn’t about replacing developers; it’s about augmenting them. Let’s walk through how you can set this up to accelerate your own testing process.

Prerequisites

Before we dive in, make sure you have the following ready:

• Python 3.8 or newer installed on your system.
• An active OpenAI API key. You can get one from the OpenAI platform dashboard.
• The necessary Python libraries. You’ll need to install them using your package manager, for instance by running `pip install openai python-dotenv`.

The Guide: Step-by-Step

Step 1: Project Setup

First, let’s get our project structure in order. I’ll skip the standard virtual environment setup, as you probably have your own workflow for that. Let’s jump straight to the files. We’ll need three things:

1. config.env: To securely store our API key. Never hardcode credentials.
2. utils.py: The Python module with the functions we want to test.
3. generate_tests.py: The script that will do the heavy lifting.

In your `config.env` file, add your API key like this:

OPENAI_API_KEY="your-super-secret-key-goes-here"

Step 2: The Code We Want to Test

Let’s create some simple functions in a file named `utils.py`. This will be our target. A mix of simple and slightly more complex logic is good for seeing how the AI handles it.

# utils.py
import re

def is_valid_email(email):
    """Checks if the provided string is a valid email format."""
    if not isinstance(email, str):
        return False
    pattern = r"^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$"
    return re.match(pattern, email) is not None

def extract_numbers_from_string(text):
    """Extracts all integer numbers from a string and returns them as a list."""
    if not isinstance(text, str):
        raise TypeError("Input must be a string.")
    
    return [int(num) for num in re.findall(r'\d+', text)]

Step 3: The Test Generation Script

Now for the magic. We’ll write `generate_tests.py`. This script will read the source code from `utils.py`, construct a precise prompt for the AI, call the OpenAI API, and save the generated test skeleton into a new file.

Here’s the complete script. I’ll break down what each part does below.

# generate_tests.py
import os
from openai import OpenAI
from dotenv import load_dotenv

def generate_test_skeleton(source_file_path):
    """
    Reads a Python file and uses OpenAI to generate a unit test skeleton for it.
    """
    # 1. Load environment variables and configure the API client
    load_dotenv('config.env')
    api_key = os.getenv("OPENAI_API_KEY")
    if not api_key:
        print("Error: OPENAI_API_KEY not found in config.env")
        return

    client = OpenAI(api_key=api_key)

    # 2. Read the source code from the target file
    try:
        with open(source_file_path, 'r') as f:
            source_code = f.read()
    except FileNotFoundError:
        print(f"Error: Source file not found at {source_file_path}")
        return

    # 3. Construct the prompt for the AI
    # This is the most important part!
    prompt = f"""
    Based on the following Python code, generate a unit test skeleton file using the 'unittest' library.

    Rules:
    1.  Create a test class named 'Test[ModuleName]' where [ModuleName] is the capitalized name of the source file.
    2.  For each function in the source code, create a corresponding test method named 'test_[function_name]'.
    3.  Include necessary imports, including the functions from the source module.
    4.  The body of each test method should contain only the 'pass' statement. Do not write any assertions.
    5.  Add a main execution block `if __name__ == '__main__': unittest.main()` at the end.
    6.  The output should be only the Python code for the test file. Do not include any explanations.

    Source Code from '{source_file_path}':
    ---
    {source_code}
    ---
    """

    # 4. Make the API call
    print(f"Generating test skeleton for {source_file_path}...")
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo",
            messages=[
                {"role": "system", "content": "You are a helpful assistant that generates Python unit test code."},
                {"role": "user", "content": prompt}
            ],
            temperature=0.2 # We want deterministic, not creative, output
        )
        generated_code = response.choices[0].message.content.strip()

        # Clean up the response to ensure it's just code
        if "```python" in generated_code:
            generated_code = generated_code.split("```python")[1].split("```")[0].strip()

    except Exception as e:
        print(f"An error occurred while calling the OpenAI API: {e}")
        return

    # 5. Save the generated skeleton to a new file
    output_filename = f"test_{os.path.basename(source_file_path)}"
    with open(output_filename, 'w') as f:
        f.write(generated_code)

    print(f"Successfully generated test skeleton at {output_filename}")


if __name__ == '__main__':
    target_file = 'utils.py'
    generate_test_skeleton(target_file)

Logic Breakdown:

• Part 1 & 2: We’re just doing some standard setup—loading the API key from our `config.env` and reading the content of `utils.py`.
• Part 3 (The Prompt): This is where the real engineering happens. Notice how specific the instructions are. We explicitly ask for the `unittest` library, define the naming convention for classes and methods, and crucially, tell it to use `pass` in the method bodies. This prevents the AI from guessing at assertions, which is where it often makes mistakes.

Pro Tip: The quality of your generated code is 100% dependent on your prompt. Be overly specific. I found that providing rules in a numbered list and using a clear separator for the source code (like `—`) dramatically improves the reliability of the output.

• Part 4: We make the API call. I’m using `gpt-3.5-turbo` because it’s fast and cheap for this kind of task. A low `temperature` (like 0.2) makes the output more focused and less random. We also do a bit of cleanup to strip out any conversational text or markdown formatting the AI might add.
• Part 5: We write the cleaned-up response to a new file, `test_utils.py`.

Step 4: Run It and See the Result

Now, just run the script from your terminal: `python3 generate_tests.py`

You should see a new file, `test_utils.py`, appear in your directory. It should look something like this:

# test_utils.py
import unittest
from utils import is_valid_email, extract_numbers_from_string

class TestUtils(unittest.TestCase):

    def test_is_valid_email(self):
        pass

    def test_extract_numbers_from_string(self):
        pass

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

Look at that! It’s a perfect, clean skeleton. All the boilerplate is done. Now a developer can go in and fill out the `pass` statements with meaningful assertions (e.g., `self.assertTrue(is_valid_email(‘test@example.com’))`).

Common Pitfalls

I’ve run this workflow dozens of times, and here’s where I usually mess up, so you don’t have to:

• Weak Prompts: My first few prompts were too vague, like “Write tests for this code.” The AI would sometimes use `pytest`, sometimes `unittest`, and often invent crazy assertions. Being explicit in the prompt is non-negotiable.
• API Key Exposure: In a moment of haste, I once hardcoded a key. Don’t do it. Always use environment variables, preferably loaded from a file that is in your `.gitignore`.
• Expecting Perfection: Remember, this is a tool for creating *skeletons*. The AI might misunderstand a complex function or miss an edge case. You still need a human to write the actual test logic. Think of it as a smart intern who does the setup for you.

Conclusion

Integrating an LLM into your development workflow like this is a massive productivity booster. It handles the monotonous, low-level tasks, freeing up your team’s brainpower for complex problem-solving. By generating test skeletons, you lower the barrier to writing tests, which can significantly improve code coverage and quality, especially on legacy projects where tests are sparse.

Give it a try. It takes about 15 minutes to set up, and I guarantee it’ll save you far more time in the long run.

– Darian Vance

Darian Vance

Lead Cloud Architect & DevOps Strategist

With over 12 years in system architecture and automation, Darian specializes in simplifying complex cloud infrastructures. An advocate for open-source solutions, he founded TechResolve to provide engineers with actionable, battle-tested troubleshooting guides and robust software alternatives.

🤖 Frequently Asked Questions