My name is Piotr, a passionate pythonista and this is my blog!

    Smart mock by side effect replacement

    Posted at — Jan 17, 2023

    Smart way of mocking a Python method using side_effect for multiple different return values with ability to preserve the input of mocked method as part of output.

    The use-case is quite simple. We have a method that we want to mock but this method returns part of input data in the output.

    class OrderManager:
        def get_order_details(self) -> Dict[str, Any]:
            return {"name": "John Smith"}
        def get_order(self, order_id: int) -> Dict[str, Any]:
            book_data = self.get_order_details()
            return {"id": order_id, **book_data}
        def get_many_orders(self, orders_ids: List[int]) -> List[Dict[str, Any]]:
            # NOTE: Some pre-get logic here
            orders = [self.get_order(order_id) for order_id in orders_ids]
            # NOTE: Some post-get logic here
            return orders

    We want to test the important logic in get_many_orders and do not care about anything else. We can assume that get_order_details method makes HTTP request and return some data.

    from unittest import mock
    def test_delete_book(mock_get_order: mock.MagicMock) -> None:
      orders_ids = [1, 2]
      class SmartMock:
          def __init__(self, *side_effects):
              self.side_effects = iter(side_effects)
          def __call__(self, id):
              side_effect = next(self.side_effects)
              return side_effect(id)
      def harry_potter_order(order_id):
          return {"name": "Harry Potter", "id": order_id}
      def frodo_order(order_id):
          return {"name": "Frodo", "id": order_id}
      mock_get_order.side_effect = SmartMock(harry_potter_order, frodo_order)
      books = OrderManager.get_many_orders(orders_ids)
      assert books == [harry_potter_order(1), frodo_order(2)]

    To do it, we create python iterator to call each side effect one by one when the tested method is executed so at the end the input data placed in the output are the same, and we do not need to manipulate it manually.