Writing Reusable and Testable Code

# test_atm.py
# importing withdrawal function from atm.py
import unittest
from atm import withdrawal

class TestATMWithdrawal(unittest.TestCase):

    def test_insufficient_funds(self):
        message, new_balance = withdrawal(100,150)
        self.assertEqual(message, "Insufficient funds")
        self.assertEqual(new_balance, 100)

    def test_invalid_amount(self):
        message, new_balance = withdrawal(100,-20)
        self.assertEqual(message, "Invalid amount")
        self.assertEqual(new_balance, 100)

    def test_successful_transaction(self):
        message, new_balance = withdrawal(100,50)
        self.assertEqual(message, "Transaction successful")
        self.assertEqual(new_balance, 50)

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

# From a different file
# atm.py
def withdrawal(balance,amount):
    if amount > balance: 
        return "Insufficient funds", balance
    elif amount <= 0:
        return "Invalid amount", balance
    else:
        balance -= amount
        return "Transaction successful", balance
    
print(withdrawal(100,150))