Byte Ebi's Logo

Byte Ebi 🍤

每天一小口，蝦米變鯨魚

[Pytest 101] 04 Mock Class

介紹兩種在使用工廠模式時測試 class 建構的方式，並說明如何有效抽換依賴。

Ray

在開發中導入工廠模式雖然能讓架構更清晰，但也讓單元測試變得複雜。
一次搞懂兩種不同 mock 策略的優劣與實作方式，幫助你寫出更穩健的測試程式！

當程式碼中使用工廠模式（Factory Design Pattern）分別建構不同類別回傳時，有兩種單元測試的寫法。

# my_project/service/alert_service.py

from service.alert.mail_alert import MailAlert
from service.alert.teams_alert import TeamsAlert
from service.alert.discord_alert import DiscordAlert

def get_alert_model(session, alert_type, alert_config):
    if alert_type == "mail":
        return MailAlert(session, alert_config)
    elif alert_type == "teams":
        return TeamsAlert(session, alert_config)
    elif alert_type == "discord":
        return DiscordAlert(session, alert_config)
    else:
        raise Exception(
            f"Alert Type is not valid. Please check the alert type mail, teams or discord. {alert_type}"
        )

Mock class

因為已經 import 到 Module 中，所以抽換的對象是要測試的檔案中的類別

  1. service.alert_service.MailAlert
  2. service.alert_service.TeamsAlert
  3. service.alert_service.DiscordAlert
# my_project/test/service/test_alert_service.py

import pytest
from pytest_mock import MockFixture

from service.alert_service import (
    get_alert_model,
)

def test_get_alert_model_success(mocker: MockFixture):
    mock_session = mocker.MagicMock()
    alert_config = {}

    alert_type_model_dict = {
        "mail": "MailAlert",
        "teams": "TeamsAlert",
        "discord": "DiscordAlert",
    }

    for alert_type, alert_model in alert_type_model_dict.items():
        mock_alert_model = mocker.patch(f"service.alert_service.{alert_model}")

        result = get_alert_model(mock_session, alert_type, alert_config)

        assert result == mock_alert_model.return_value
        mock_alert_model.assert_called_once_with(mock_session, alert_config)


def test_get_alert_model_invalid_type(mocker: MockFixture):
    mock_session = mocker.MagicMock()

    alert_type = "test_type"
    with pytest.raises(
        Exception,
        match=f"Alert Type is not valid. Please check the alert type mail, teams or discord. {alert_type}",
    ):
        get_alert_model(mock_session, alert_type, {})

說明

先使用 mocker.patch 將對應的類別 mock 出來成 mock_alert_model 物件
再用斷言式判斷函式得到的結果與 mock_alert_model 相同
最後使用 assert_called_once_with 確認 mock_alert_model 只被呼叫一次，且該次呼叫使用了指定的引數

更多可用的斷言方法：使用 Mock 的方式

引入實際類別

另一種方法是在測試中 import 類別對真實類別產生依賴。
這樣做可以更貼近實際執行的狀況，但是測試的撰寫會更加複雜。

此方法可以使用 isinstance 來判斷類別，當引入的類別被修改時可能會對測試案例產生影響。
並且如果傳入的測試資料要無法建立對應的類別，會直接報錯。
以這邊來說就是如果建構類別時會使用到 alert_config 中的參數 alert_content
而我們傳入的是空 dict 的時候就會發生錯誤：

AttributeError: ‘dict’ object has no attribute ‘alert_content’

# my_project/test/service/test_alert_service.py

import pytest
from pytest_mock import MockFixture

from service.monitor.alert.mop_mail_alert import MailAlert
from service.monitor.alert.teams_alert import TeamsAlert
from service.monitor.alert.trigger_airflow_dags import DiscordAlert
from service.alert_service import (
    get_alert_model,
)

def test_get_alert_model_success(mocker: MockFixture):
    mock_session = mocker.MagicMock()
    alert_config = {}

    # 使用真實類別
    alert_type_model_dict = {
        "mail": MailAlert,
        "teams": TeamsAlert,
        "discord": DiscordAlert,
    }

    # 不對類別進行 mock
    for alert_type, alert_model in alert_type_model_dict.items():
        result = get_alert_model(mock_session, alert_type, alert_config)
        assert isinstance(result, alert_model)

最新文章

Category

Tag