python - 如何模拟自定义 throttle 类以在 django-rest-framework 中进行测试？

Question

我有一个自定义 throttle 类，如:(打印语句用于调试 :) ) 在 api.throttle.py

print("befor CustomThrottle class")
class CustomThrottle(BaseThrottle):
    def __init__(self):
        super().__init__()
        print("initializing CustomThrottle", self)
        self._wait = 0

    def allow_request(self, request, view):
        print("CustomThrottle.allow_request")
        if request.method in SAFE_METHODS:
            return True
        # some checking here
        if wait > 0:
            self._wait = wait
            return False
        return True

    def wait(self):
        return self._wait

我的 api.views.py就好像:

from api.throttle import CustomThrottle

print("views module")
class SomeView(APIView):
    print("in view")
    throttle_classes = [CustomThrottle]

    def post(self, request, should_exist):
        # some processing
        return Response({"message": "Done."})

我的测试是 api/tests/test_views.py :

    @patch.object(api.views.CustomThrottle, "allow_request")
    def test_can_get_confirmation_code_for_registered_user(self, throttle):
        throttle.return_value = True
        response = self.client.post(path, data=data)
        self.assertEqual(
            response.status_code,
            status.HTTP_200_OK,
            "Should be successful",
        )

    @patch("api.views.CustomThrottle")
    def test_learn(self, throttle):
        throttle.return_value.allow_request.return_value = True
        response = self.client.post(path, data=data)
        self.assertEqual(
            response.status_code,
            status.HTTP_200_OK,
            "Should be successful",
        )

第一个测试正确通过，但仍然是 CustomThrottle类已实例化，但 allow_request方法被模拟；第二个测试表明 CustomThrottle类被 mock ，也不是 allow_request方法，它失败了 status 429 ( CustomThrottle 的节流率为 2 分钟)。

Answer 1

如果您还想模拟 CustomThrottle类，以及他的属性，你还需要patch在您需要它的地方，因为 patch装修申请monkey patching在你调用它的地方。

关于您的情况，您可以执行以下操作

from unittest.mock import MagicMock, patch

    @patch('api.views.CustomThrottle')
    @patch('api.views.PhoneConfirmationThrottle')
    def test_learn(self, phone_throttle, custom_throttle):
        mocked_custom_throttle_instance = MagicMock()
        mocked_custom_throttle_instance.allow_request = True
        mocked_custom_throttle_instance.status = ...

        # Do some stuff which do you need with the mocked class
        # and then return an instance for your patched object

        custom_throttle.return_value = mocked_custom_throttle_instance

        response = self.client.post(path, data=data)
        self.assertEqual(
            response.status_code,
            status.HTTP_200_OK,
            "Should be successful",
        )

它将取代您的 real对象与 mocked一，还退房 unittest.mock documentation ，它会在 future 帮助你。