Skip to main content
In this article, we enhance our BankAccount class by introducing exception handling for withdrawal operations. This update ensures that users cannot withdraw more money than is available in their account. The guide covers the original implementation, the issue with insufficient balance checks, testing adjustments, and the introduction of a custom exception for better error handling.

Original BankAccount Implementation

Initially, our BankAccount class was implemented as follows:
When running tests, all tests passed:

Problem: Insufficient Funds Check

Upon inspection, the withdraw method does not verify if the account has sufficient funds. For example, if the account has 100andawithdrawalof100 and a withdrawal of 500 is attempted, the operation should be blocked. To address this, we add a check in the withdraw method that raises an exception when necessary.

Updated BankAccount Implementation with Basic Exception

The revised BankAccount class adds a condition to the withdraw method to raise an exception if the withdrawal amount exceeds the account balance:
Re-running the tests produces the following output:

Adjusting Tests for Exception Handling

The original parameterized test included a case where more money was withdrawn than deposited:
The test output indicated a failure for the case where 10 is deposited and 200 is withdrawn:
Since an exception is expected for insufficient funds, it’s better to remove this failing test case from the parameterized group and create a dedicated test.

Creating a Dedicated Test for Insufficient Funds

We use Pytest’s raises context manager for this purpose. For an account with an initial balance of zero (using the zero_bank_account fixture), attempting to withdraw money should trigger an exception:
Here, the bank_account fixture initializes an account with a balance (for instance, 50) to ensure that withdrawing $200 raises the expected exception. Running the tests now produces:

Introducing a Custom Exception: InsufficientFunds

For better clarity and maintainability in production code, we define a custom exception called InsufficientFunds that inherits from Python’s built-in Exception class. This approach allows our tests to verify that the correct exception is raised.

Defining the Custom Exception

Updated BankAccount Class with Custom Exception

The BankAccount class is updated to raise InsufficientFunds when an over-withdrawal is attempted:

Updating the Test for Insufficient Funds

We also update the insufficient funds test to assert that an InsufficientFunds exception is raised:
Running the tests now confirms that the proper exception is raised:
Specifying the exact exception type ensures that our code not only raises an error but also raises the correct, expected type. This improves test precision and code reliability.

Demonstrating the Importance of the Correct Exception Type

To highlight the importance of using a custom exception, consider a scenario where ZeroDivisionError is raised instead of InsufficientFunds. For example, changing the withdraw method as shown below would trigger a test failure:
Test output in this case would be:
This failure confirms that our tests correctly check for a specific exception, and any deviation from the expected behavior will be flagged immediately.

Conclusion

By defining and raising a custom exception for insufficient funds, our BankAccount class now operates in a robust and predictable manner. The updated tests validate this behavior precisely, ensuring that the internal logic of our application remains stable and error-free. Happy testing!

Watch Video