Generate Unit Tests for Python Code with watsonx Code Assistant

In modern software development, writing Python code is only the first step. Ensuring correctness, reliability and maintainability require thorough unit testing.

Python provides built-in tools such as the unittest module, allowing developers to import unittest and structure test.py files effectively. For complex or enterprise-level applications, writing unit tests manually can be time-consuming and prone to human error. Watsonx Code Assistant® (WCA) further enhances this workflow by helping developers write unit tests automatically, creating unittest.TestCase classes and generating test suites to cover both normal and edge scenarios.

Unit testing plays a critical role in validating Python functions, methods and individual pieces of code by ensuring they behave as expected under different scenarios.

Reliability: Unit tests help detect failures early in the development cycle, preventing unexpected runtime errors. By validating expected outputs in advance, unit testing ensures predictable and stable application behavior before code is deployed.
Debugging: Unit testing makes it easier to identify common runtime issues:

AssertionError occurs when a test’s expected outcome does not match the actual result returned by the code, typically indicating a logic or implementation issue.
TypeError occurs when an operation is performed on incompatible data types, such as passing a string where a numeric value is expected.
Other runtime exceptions (for example, ValueError, IndexError) indicate invalid inputs or incorrect assumptions in the code.

By isolating these errors at the unit level, developers can quickly locate the source of the problem and correct it with minimal impact.

Maintainability: Unit tests enable safe refactoring by ensuring that existing functionality continues to work as intended. When changes are made, tests act as a safety net that helps prevent regressions and unintended side effects.
Automation: Unit testing supports automated execution by using test runners and command-line tools, making it easy to integrate tests into CI/CD pipelines. This method ensures continuous validation of code changes without manual intervention.

Consistent unit testing promotes reusable and well-structured code, standardized test modules and reliable CI/CD integration. Overall, it reduces defects, improves Python code quality and enhances development productivity.

Prompt given to WCA:

Generate unit tests by using Python unittest framework for the following function for normal cases, edge cases and error scenarios.

#Source code provided to WCA
def divide(a, b):
    if b == 0:
        raise ValueError(“Division by zero is not allowed”)
    return a / b

Powered by Granite

#Code generated by WCA
import unittest
from calculator import divide
class TestDivide(unittest.TestCase):
    def test_valid_division(self):
        self.aself.assertEqual(divide(10, 2), 5)
    def test_division_result_float(self):
        self.assertEqual(divide(7, 2), 3.5)
    def test_division_by_zero(self):
        with self.assertRaises(ValueError):
            divide(5, 0)
# Required for jupyter to ignore kernel arguments
if __name__ == “__main__”:
    unittest.main(argv=[‘first-arg-is-ignored’], exit=False)

Powered by Granite

Before we begin to use watsonx Code Assistant (WCA), ensure that the following prerequisites are completed.

Step 1: Install Visual Studio Code

Watsonx Code Assistant is available as an extension for Visual Studio Code.

Install the latest stable version of Visual Studio Code.
Verify that Visual Studio Code launches correctly on our system.

Step 2: Install Python

For Python development and unit test generation,

Install Python version 3.9 or later
Verify the installation by using python --version
Ensure that pip is available
Create and activate a Python virtual environment

Step 3: Install the watsonx Code Assistant extension

Open Visual Studio Code.
Go to view > extensions.
Search for watsonx Code Assistant.
Install the extension.

Use this link https://w3.ibm.com/w3publisher/wca-ibm/download/vscode for installation details and once installed, the extension is available within Visual Studio Code.

Step 4: Choose a backend (cloud or local)

Before configuring watsonx Code Assistant, it is important to understand the differences between cloud and local backends.

Why choose cloud versus local?

Cloud backend is recommended for individual users and quick setup, providing managed infrastructure and minimal configuration.
Local backend is typically used in enterprise environments where data privacy, compliance or internal infrastructure requirements apply.

4.1 Cloud backend

Uses a managed backend service
Requires minimal setup
Suitable for external users and personal development
Automatically handles model updates and scaling

4.2 Local backend

Runs within local or enterprise-managed infrastructure
Backend configuration is managed by the organization
Commonly used in regulated or security-sensitive environments

After selecting the wanted backend, restart Visual Studio Code to ensure that the configuration is applied successfully.

Step 5: Configure watsonx Code Assistant

Open Visual Studio Code.
Navigate to File > Preferences > Settings > Extensions > watsonx Code Assistant.
Select the appropriate backend (cloud or local).
Save the settings.
Restart Visual Studio Code.

Step 6: Create or open a project

Open an existing project or create a new project folder.
Ensure Python source files (.py) are visible in the workspace.
Activate a virtual environment if applicable.

Step 7: Verify installation

Open a Python file in Visual Studio Code.
Select a function or class.
Use watsonx Code Assistant features such as unit test generation or code explanation.

If the assistant responds successfully, the setup is complete.

Note for enterprise readers: Performance and data privacy

Performance: Cloud backends offer scalable performance with minimal setup, while local backends depend on internal infrastructure capacity.
Data privacy: Local backends help ensure that source code and prompts remain within organizational boundaries, making them suitable for compliance-driven environments.

Watsonx Code Assistant (WCA) relies on VS Code editor APIs to provide context aware features such as code explanation, refactoring and unit test generation. These APIs are available for standard source code files like .py, .java and .js.

Jupyter Notebooks (.ipynb) use a kernel-based execution model and a notebook-specific interface that does not expose the same editor-level APIs. Because of this architectural difference, WCA cannot attach to or operate within individual notebook cells. As a result, the assistant is not visible in .ipynb files and this behavior is expected.

Developer tip:
For data science workflows, use WCA in .py files to generate or modify code and unit tests, then execute or validate that logic interactively in Jupyter Notebooks.

Watsonx Code Assistant (WCA) automates the creation of Python unit tests. The generated tests are executed by using Python’s unittest framework, providing a clear separation between test generation and test execution within the testing lifecycle.

Step-by-step workflow to generate and run unit tests

Step 1: Open a Python source file

Open a .py file in VS Code that contains one or more Python functions (for example, add() or divide()).

Step 2: Select the target function

Highlight the function for which you want to generate unit tests.

Step 3: Generate unit tests by using WCA

From the editor or command palette, select watsonx Code Assistant → Generate unit tests.

Step 4: Review the generated test file

WCA analyses the selected function and generates a test file (for example, test_calculator.py) that includes

A unittest.TestCase class
Multiple test methods
Assertions for normal, edge and error scenarios

No manual test creation is required.

Step 5: Run the generated tests

Run the generated tests by using Python’s unittest framework and python -m unittest. Python is responsible for executing the tests, while WCA is responsible for generating them.

1. Functions

Functions are the core building blocks of python code. Each function should be validated with unit tests to ensure that it produces the correct output for valid inputs, handles edge cases and raises appropriate exceptions when necessary. Testing functions is essential in both small scripts and large codebases, where multiple pieces of code interact.

2. Testcase class

A testcase class is a subclass of unittest.Testcase and serves to group related test methods. Testcase classes help organize test code, making it readable, maintainable and modular. Each test method within a testcase class typically tests a single function or a specific scenario.

3. Assertions

Assertions are used to verify that code behaves as expected. Python’s unittest module provides multiple assert methods:

assertequal: Verifies equality between expected and actual values.
assertfalse: Ensures that a condition is false.
assertin: Checks whether an element exists within a collection.
assertis: Confirms that two objects are the same instance.
assertraises: Ensures that a specific exception is raised, useful for error handling and validation of typeerror or assertionerror scenarios.

Well-crafted assert statements allow developers to catch failed tests early and improve the reliability of the Python codebase.

4. Test fixtures

Test fixtures manage setup and teardown routines before and after test runs:

setupclass: Prepares the resources needed for all tests in a testcase class.
teardown: Cleans up resources after each test or after the entire suite runs.

Fixtures ensure tests run in a controlled and consistent environment, preventing side effects between test functions and reducing debugging complexity.

5. Test suite and test runner

A test suite is a collection of testcase classes or individual test functions, organized to run together. The test runner executes the suite, producing a verbose output to highlight test passes and failed tests, including detailed traceback logs.

unittest.main: Provides a simple command-line interface to run test suites automatically.
Test discovery: Automatically finds test.py files and test functions across the codebase, streamlining unit testing for large projects.
Using test suites and test runners, developers can maintain comprehensive test coverage for all functions, subclasses and pieces of code, ensuring correctness across the Python codebase.

import unittest
from calculator import add, divide

# 1. Functions example
class TestFunctions(unittest.TestCase):
    def test_add_function(self):
        self.assertEqual(add(2, 3), 5)

    def test_divide_function(self):
        self.assertEqual(divide(10, 2), 5)
        with self.assertRaises(ValueError):
            divide(10, 0)
# 2. Testcase class example
class TestMath(unittest.TestCase):
    def test_basic(self):
        self.assertTrue(10 > 5)
# 3. Assertions examples
class TestAssertions(unittest.TestCase):
    def test_assertions(self):
        self.assertEqual(add(2, 2), 4)
        self.assertFalse(False)
        self.assertIn(3, [1, 2, 3])
        self.assertIs(None, None)
    with self.assertRaises(ValueError):
        divide(10, 0)
# 4. Test fixtures example
class TestUser(unittest.TestCase):
@classmethod
def setUpClass(cls):
    cls.sample_data = {“name”: “John”, “age”: 30}
def test_name(self):
    self.assertEqual(self.sample_data[“name”], “John”)
@classmethod
def tearDownClass(cls):
    cls.sample_data = None
# 5. Test suite and runner example
class TestSuiteExample(unittest.TestCase):
    def test_add_suite(self):
        self.self.assertEqual(add(5, 5), 10)
def suite():
    suite = unittest.TestSuite()
    suite.addTest(TestFunctions(“test_add_function”))
    suite.addTest(TestFunctions(“test_divide_function”))
    suite.addTest(TestMath(“test_basic”))
    suite.addTest(TestAssertions(“test_assertions”))
    suite.addTest(TestUser(“test_name”))
    suite.addTest(TestSuiteExample(“test_add_suite”))
    return suite
if __name__ == “__main__”:
    runner = unittest.TextTestRunner()
    runner.run(suite())

Powered by Granite

The two most commonly used frameworks are unittest and pytest, each suited to different development needs.

1. unittest module

The unittest module is Python’s built-in framework for unit testing and is part of the standard library.

Key features include:

Creation of unittest.TestCase classes
Organization of test methods in test files
Automated test execution by using unittest.main()
Built-in assertion methods for validating return values and exception handling

When to use unittest:
unittest is well suited for enterprise and large-scale projects that require a structured, standardized and widely supported testing approach. It provides a clear separation between production code and test code, making it easier to maintain and govern large Python codebases.

2. pytest framework

pytest is a popular third-party testing framework known for its simplicity and flexibility.

Key features include:

Lightweight setup
Automatic discovery of test files and test functions
Support for plug-ins, decorators and detailed output
Effective for testing shared logic across multiple modules

When to use pytest:
pytest is ideal for rapid prototyping, data-driven testing and experimentation, particularly in data science and analytics workflows where fast iteration and parameterized tests are common.

Summary: Choosing the right framework

Use unittest for enterprise-grade, structured and compliance-driven projects
Use pytest for rapid development, experimentation and data-centric testing scenarios

Watsonx Code Assistant (WCA) provides a major boost in efficiency for writing unit tests:

Analyses the Python code and identifies pieces of code suitable for testing.
Automatically generates unittest.TestCase subclasses, test functions and assert statements.
Builds comprehensive test suites covering normal, edge and error scenarios.
Supports both unittest and pytest frameworks for flexibility.
Reduces debugging time and ensures higher validation coverage for our source code.

Automation with WCA ensures that even large codebases are systematically tested without requiring manual effort, allowing developers to focus on software development rather than writing repetitive tests.

Organize test.py files logically, reflecting the codebase structure.
Write clear and descriptive testcase classes for readability.
Cover both positive and negative scenarios, including typeerror and assertionerror cases.
Use setupclass and teardown routines for reusable test fixtures.
Run tests automatically by using unittest.main or test runners from the CLI.
Integrate tests with GitHub to automate CI/CD pipelines.
Regularly refactor test modules to maintain quality and reliability.
Use assert methods strategically to catch subtle bugs in pieces of code.
Leverage decorators in pytest for parameterized or conditional tests.

Following these best practices ensures that the Python testing workflow is reliable, repeatable and maintainable across teams and projects.

Even with automation, understanding debugging is essential for high-quality Python code:

Use verbose output to monitor failed tests.
Examine traceback logs to pinpoint the origin of exceptions.
Apply assert statements consistently to validate return values and function behaviors.
Ensure pieces of code reused across py files are adequately tested.
Run test discovery regularly to catch untested functions or modules.

Proper debugging complements unit testing, helping maintain reliability across Python projects.

Test discovery: Automatically detects test.py files and test functions across large codebases.
Subclasses: Advanced functions can be tested in unittest classes by using subclassing techniques.
Decorators: Parametrize or skip test methods conditionally.
Plug-ins: Extend pytest for coverage, reporting or mocking external dependencies.
Cross-platform testing: Run unit tests on Linux, Windows and cloud environments.
GitHub integration: Trigger test runs automatically on commits and pull requests.
Incremental testing: Continuously write unit tests as new functions and pieces of code are added.

These advanced strategies enhance coverage, reduce failed tests and simplify debugging, ensuring robust and scalable software development practices.

Unit tests can be applied in diverse scenarios:

Function validation: Ensuring functions return expected results for all inputs.
Error handling: Capturing assertionerror, typeerror and other exceptions.
API testing: Validating endpoints in data science or enterprise APIs.
Refactoring: Ensuring unittest classes detect regressions during source code refactors.
Edge cases: Testing subclasses, decorators or complex Python functions.
Automation: Running test suites through unittest.main or test runners to validate large codebases.

Automated unit test generation with WCA ensures comprehensive coverage and reduces manual effort in all these use cases.

Integrating unit tests into CI/CD pipelines (for example, with GitHub) ensures that every test file, test case and test method is automatically validated as part of the development workflow.

Key benefits include

Commit and push test files to the source control repository
Automatically trigger test runners on commits or pull requests
Monitor test execution results and verbose output for failures
Maintain test suites that cover functions, classes and subclasses
Reduce manual debugging and improve overall Python code quality across distributed teams

This integration transforms unit testing from a manual activity into an automated, repeatable and scalable process.

For enterprise environments, integrating unit tests into CI/CD pipelines also supports:

Auditability: Test execution results are logged and traceable for compliance reviews
Consistency: Standardized test execution ensures uniform quality checks across teams
Policy enforcement: Pipelines can enforce quality gates, preventing code merges when tests fail
Security and reliability: Early detection of defects reduces production risk and supports regulatory requirements

By embedding unit testing into CI/CD pipelines, organizations can align software quality practices with enterprise governance, compliance and operational standards.

Watsonx Code Assistant empowers developers to streamline and standardize Python unit testing by enabling:

Efficient generation of unit tests for functions, classes and subclasses
Automated creation of unittest.TestCase classes in test files
Comprehensive test coverage for normal scenarios, edge cases and exceptions
Seamless integration of unit tests with GitHub and CI/CD pipelines
Reduced debugging effort and improved overall Python code quality

By leveraging WCA, developers can focus more on application logic while maintaining reliable, maintainable and scalable test suites across both small and enterprise projects.

Developer tips and best practices

Keep unit tests in a separate tests/ directory to improve maintainability and project structure.
Use WCA to quickly bootstrap unit tests, then refine assertions manually to reflect business logic and domain-specific rules.
Run unit tests locally before pushing changes to ensure faster feedback and cleaner CI/CD runs.

Generate unit tests for Python code with watsonx Code Assistant

Introduction

Why unit testing is essential?

Demonstrating unit test generation by using watsonx Code Assistant (WCA)

Prerequisites

Why is watsonx Code Assistant available in .py files but not in Jupyter Notebooks?

How does WCA generate unit tests?

Key components of Python unit testing

Python testing frameworks

Automating unit test generation with watsonx Code Assistant

Best practices for unit testing

Debugging and validation

Advanced unit testing concepts

Unit testing use cases

Integrating unit tests in CI/CD pipelines

Enterprise compliance and governance considerations

Conclusion

Author

Resources