119. 'Capstone: Unittest.Mock + Enum + Typing'
A typed API client with status enums, Optional return types, and mock-based tests
119. βCapstone: Unittest.Mock + Enum + Typingβ
π A typed API client with status enums is tested with mocks. The developer compares an Enum member to a plain string instead of using the enum value, so the condition never triggers. They also forget to handle the Optional return β calling .temp on None crashes.
π‘ Fun fact: The pattern of using Enum for API status codes (SUCCESS, ERROR, NOT_FOUND) combined with Optional return types is a Python implementation of a concept from functional programming called a βResult typeβ or βOption type.β Languages like Rust (Result<T, E>), Haskell (Maybe), and Swift (Optional) enforce this at the type system level. In Python, itβs conventional β type checkers like mypy will warn you when you access attributes on an Optional[X] without first checking for None, but Python itself wonβt stop you.
β οΈ Watch out: ResponseStatus.SUCCESS == "success" is always False. An Enum member is a Python object of type ResponseStatus, and a plain string is of type str β they can never be equal. When you receive status values from an external API (which returns raw strings in JSON), you must compare using .value: raw_string == ResponseStatus.SUCCESS.value. Or construct the Enum member from the string: ResponseStatus(raw_string) == ResponseStatus.SUCCESS.
π€ Think about it: The get_feels_like() method returns float but calls get_weather() which returns Optional[WeatherResponse]. Without a None guard, calling response.feels_like when response is None raises AttributeError β a runtime crash. Could a type checker like mypy have caught this bug statically? What does the return type annotation of get_feels_like() -> float promise, and does the broken code fulfill that contract?
Learning objectives
- Compare Enum members to strings using .value, not the member directly
- Guard Optional return values with an explicit None check before attribute access
- Mock instance methods by assigning a MagicMock to the attribute
- Use side_effect to make a mock raise an exception
- Design typed APIs with Enum status codes and Optional returns for safe callers
Key concepts
- Enum.MEMBER.value β the underlying string/int value
- Enum.MEMBER == plain_string is always False
- Optional[X] requires a None guard before attribute access
- mock.side_effect = Exception() β make mock raise
- client.method = MagicMock(β¦) β replace instance method
Try it
Concept detail
Capstone: unittest.mock + enum + typing
Enum β Named Constants with Values
from enum import Enum
class Status(Enum):
SUCCESS = "success"
ERROR = "error"
# Compare members to members:
if response.status == Status.SUCCESS: # correct
# Compare .value to strings (from JSON/API):
if raw_string == Status.SUCCESS.value: # correct
# WRONG β Enum member != plain string:
if raw_string == Status.SUCCESS: # always False!Optional β Guard Against None
from typing import Optional
def get_data() -> Optional[dict]:
return None # might happen
result = get_data()
if result is None: # ALWAYS check before using
return default_value
print(result["key"]) # safe nowMocking Instance Methods
from unittest.mock import MagicMock
client = MyClient()
# Replace the method directly on the instance:
client._fetch = MagicMock(return_value={"status": "success", "data": {}})
# Or simulate failure:
client._fetch = MagicMock(side_effect=RuntimeError("network down"))Putting It Together
# Enum for status values
# Optional[WeatherResponse] for nullable return
# MagicMock to replace _fetch in tests
# .value for comparing enum to external stringsSolution
from enum import Enum
from typing import Optional, Dict, Any
from unittest.mock import MagicMock
class WeatherCondition(Enum):
SUNNY = "sunny"
CLOUDY = "cloudy"
RAINY = "rainy"
STORMY = "stormy"
UNKNOWN = "unknown"
class ResponseStatus(Enum):
SUCCESS = "success"
NOT_FOUND = "not_found"
ERROR = "error"
class WeatherResponse:
def __init__(self, status: ResponseStatus, data: Optional[Dict[str, Any]] = None):
self.status = status
self.data = data or {}
@property
def temperature(self) -> Optional[float]:
return self.data.get("temperature")
@property
def condition(self) -> WeatherCondition:
raw = self.data.get("condition", "unknown")
try:
return WeatherCondition(raw)
except ValueError:
return WeatherCondition.UNKNOWN
@property
def feels_like(self) -> Optional[float]:
return self.data.get("feels_like")
class WeatherClient:
"""Fetches weather data. In tests, _fetch is replaced by a mock."""
def _fetch(self, city: str) -> dict:
"""Internal method β replaced by mock in tests."""
raise RuntimeError("Real network call β replace with mock in tests")
def get_weather(self, city: str) -> Optional[WeatherResponse]:
"""Fetch weather for city. Returns WeatherResponse or None on failure."""
try:
raw = self._fetch(city)
except Exception:
return None
# FIX 1: Compare to the enum's .value string, or compare using the enum member
if raw.get("status") == ResponseStatus.SUCCESS.value:
return WeatherResponse(ResponseStatus.SUCCESS, raw.get("data"))
else:
return WeatherResponse(ResponseStatus.ERROR)
def get_temperature(self, city: str) -> Optional[float]:
"""Return just the temperature, or None if unavailable."""
response = self.get_weather(city)
if response is None:
return None
return response.temperature
def get_feels_like(self, city: str) -> float:
"""Return feels-like temperature. Falls back to actual temp if missing."""
response = self.get_weather(city)
# FIX 2: Guard against None response
if response is None:
return 0.0
feels = response.feels_like
if feels is None:
return response.temperature or 0.0
return feelsTests
from unittest.mock import MagicMock
def make_success_payload(temperature=22.0, condition="sunny", feels_like=20.0):
return {
"status": "success",
"data": {
"temperature": temperature,
"condition": condition,
"feels_like": feels_like,
}
}
def make_error_payload():
return {"status": "error"}
def test_enum_value_comparison():
"""ResponseStatus.SUCCESS.value should equal the string 'success'."""
assert ResponseStatus.SUCCESS.value == "success"
def test_enum_member_not_equal_to_string():
"""An Enum member is NOT equal to a plain string β this is the bug."""
# The broken code does: raw["status"] == ResponseStatus.SUCCESS
# But raw["status"] is the string "success", not the enum member
assert (ResponseStatus.SUCCESS == "success") is False, (
"Enum member equals string check should be False β compare .value instead"
)
def test_get_weather_returns_response_on_success():
"""With a success payload, get_weather should return a WeatherResponse."""
client = WeatherClient()
client._fetch = MagicMock(return_value=make_success_payload(25.0))
response = client.get_weather("Mumbai")
assert response is not None, (
"get_weather returned None on success β is the status comparison correct? "
"Use ResponseStatus.SUCCESS.value not ResponseStatus.SUCCESS"
)
assert response.status == ResponseStatus.SUCCESS
def test_get_weather_temperature_correct():
client = WeatherClient()
client._fetch = MagicMock(return_value=make_success_payload(30.0))
response = client.get_weather("Delhi")
assert response is not None
assert response.temperature == 30.0
def test_get_weather_condition_enum():
client = WeatherClient()
client._fetch = MagicMock(return_value=make_success_payload(condition="rainy"))
response = client.get_weather("Mumbai")
assert response is not None
assert response.condition == WeatherCondition.RAINY
def test_get_temperature_returns_float():
client = WeatherClient()
client._fetch = MagicMock(return_value=make_success_payload(18.5))
temp = client.get_temperature("Chennai")
assert temp == 18.5
def test_get_feels_like_does_not_crash_on_none_response():
"""get_feels_like must handle None response without AttributeError."""
client = WeatherClient()
# Make _fetch raise an exception to force None response
client._fetch = MagicMock(side_effect=RuntimeError("network down"))
try:
result = client.get_feels_like("Nowhere")
assert isinstance(result, (int, float)), (
"get_feels_like should return a number even when response is None"
)
except AttributeError as e:
assert False, (
f"get_feels_like crashed with AttributeError: {e} β "
"check for None before accessing response.feels_like"
)
def test_get_feels_like_returns_correct_value():
client = WeatherClient()
client._fetch = MagicMock(return_value=make_success_payload(22.0, feels_like=19.5))
result = client.get_feels_like("Pune")
assert result == 19.5
def test_get_weather_error_status_returns_error_response():
client = WeatherClient()
client._fetch = MagicMock(return_value=make_error_payload())
response = client.get_weather("Unknown")
assert response is not None
assert response.status == ResponseStatus.ERROR