Testing in Python for Beginners. Using `unittest` and `pytest` with Fun Examples
· 2 min read
Writing tests in Python helps ensure that your code works correctly. In this guide, we'll use two popular testing tools: the built-in unittest
module and the third-party library pytest
. We'll walk through both using examples related to vegetables and AI model names.
What Is Testing?
Testing means checking that your code behaves as expected. Instead of manually running your program over and over, you write automated tests that run every time you make a change.
unittest
— Python’s Built-in Testing Tool
Let’s start with a simple example using vegetables.
Example: Testing a Function That Capitalizes Vegetables
# veggie_tools.py
def capitalize_veggie(veggie):
return veggie.capitalize()
Now, test it using unittest
:
# test_veggie_tools.py
import unittest
from veggie_tools import capitalize_veggie
class TestVeggieTools(unittest.TestCase):
def test_capitalize_veggie(self):
self.assertEqual(capitalize_veggie("carrot"), "Carrot")
self.assertEqual(capitalize_veggie("tomato"), "Tomato")
if __name__ == "__main__":
unittest.main()
Run the test:
python test_veggie_tools.py
pytest
— A More Powerful, Cleaner Tool
pytest
makes your tests shorter and easier to read.
# test_veggie_tools_pytest.py
from veggie_tools import capitalize_veggie
def test_capitalize_veggie():
assert capitalize_veggie("cabbage") == "Cabbage"
assert capitalize_veggie("spinach") == "Spinach"
Run with:
pytest test_veggie_tools_pytest.py
✅ pytest
is simple: no need to write test classes.
Another Example: AI Model Validator
Suppose you want to validate AI model names:
# ai_model.py
def is_valid_model(name):
return name in ["GPT-4", "Claude-3", "Gemini", "LLaMA"]
unittest
Version
# test_ai_model.py
import unittest
from ai_model import is_valid_model
class TestAIModel(unittest.TestCase):
def test_is_valid_model(self):
self.assertTrue(is_valid_model("GPT-4"))
self.assertFalse(is_valid_model("GPT-2"))
pytest
Version
# test_ai_model_pytest.py
from ai_model import is_valid_model
def test_is_valid_model():
assert is_valid_model("Claude-3")
assert not is_valid_model("BERT")
When to Use What
Tool | Pros | Use if... |
---|---|---|
unittest | Comes with Python, mature | You don’t want extra dependencies |
pytest | Cleaner syntax, powerful features | You want to write lots of tests fast |
Summary
- Start with
unittest
— it’s simple and built-in. - Switch to
pytest
when you want more powerful features and cleaner tests. - Use real-world examples (vegetables, AI models) to make testing more fun!