← Home

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 now

Mocking 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 strings

Solution

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 feels

Tests

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

Resources