python:unittest中的mock模块

mock模块是为了模拟某些对象，来完成一些功能的测试。
比如，以下这个场景，你的代码，需要依赖第三方接口返回的状态码进行测试。如果第三方接口还没有完成开发，这时，你就需要使用mock模块，来模拟第三方接口的状态码。话不多说，直接上例子。

我们有一个支付类，auth是调用第三方接口的方法，调用之后，返回状态码。

• 200代表支付成功
• 500代表支付异常

这里需要实现一个支付方法pay，它依赖auth返回的状态码。

import requests


class Payment:
    """
    定义第三方支付类
    """
    @staticmethod
    def auth(card_num, amount):
        """
        请求第三方支付接口
        return：返回状态码，如果200，代表支付成功， 500代表支付异常
        """
        url = "http://第三方接口/site"
        data = {"card_num": card_num, "amount": amount}
        res = requests.post(url=url, data=data)
        return res.status_code

    def pay(self, user_id, card_num, amount):
        """
        支付方法
        """
        try:
            status_code = self.auth(card_num, amount)
        except TimeoutError:
            status_code = self.auth(card_num, amount)

        if status_code == 200:
            print(f"【{user_id}】支付【{amount}】成功！！！进行扣款并登记支付记录！")
            return "success"
        elif status_code == 500:
            print(f"【{user_id}】支付【{amount}】失败！！！不进行扣款！")
            return "Fail"

测试支付功能，写了三个场景，正常支付，异常支付，超时后又正常支付。那么通过mock函数，模拟mock方法返回的状态码，具体见下面代码：

import unittest
from unittest import mock
from scripts.handle_payment import Payment


class PaymentTest(unittest.TestCase):
    """
    测试支付接口
    """
    def setUp(self):
        self.payment = Payment()

    def test_01_success(self):
        """
        测试支付成功
        """
        self.payment.auth = mock.Mock(return_value=200)  # 
        res = self.payment.pay(user_id=10001,
                               card_num=666,
                               amount=100)

        self.assertEqual('success', res)

    def test_02_fail(self):
        """
        测试支付失败
        """
        self.payment.auth = mock.Mock(return_value=500)
        res = self.payment.pay(user_id=10002,
                               card_num=666,
                               amount=1002)

        self.assertEqual('Fail', res)

    def test_03_retry_success(self):
        """
        调用第三方接口超时之后，再次支付成功
        """
        # 重新赋值给auth的返回值，赋值为200
        self.payment.auth = mock.Mock(side_effect=[TimeoutError, 200])
        res = self.payment.pay(user_id=10003,
                               card_num=666,
                               amount=1003)

        self.assertEqual('success', res)


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