如何用不同参数Mock Python函数并断言对应返回值？

Hey there! Let's break this down step by step since you're new to mocking in Python—no worries, it's totally normal to get confused at first. Let's start by fixing your current test code, then move on to handling multiple argument-return value pairs.

First: Fix Your Existing Test Code #

Your current test has a couple of issues with order and missing steps:

1. You're calling assert_called_once_with before the mock is ever triggered (you haven't run any code that calls check_item yet).
2. You set return_value after the assertion, which means the mock won't use that value when the function is called.

Here's how to fix it, assuming you have a business function that actually calls check_item (let's say it's called process_fruit in my_module):

import unittest
import my_module
from unittest.mock import patch

class TestCheckItem(unittest.TestCase):
    @patch("my_module.check_item")
    def test_check_item_single_arg(self, mock_check_item):
        # 1. First, configure the mock's return value
        mock_check_item.return_value = 'processed_apple'
        
        # 2. Run the code that triggers the mocked function
        result = my_module.process_fruit("apple")
        
        # 3. Assert the result is correct
        self.assertEqual(result, 'processed_apple')
        
        # 4. Finally, assert the mock was called correctly
        mock_check_item.assert_called_once_with("apple")

Handling Multiple Arguments with Different Returns #

To make check_item return different values based on its input, you'll use side_effect instead of return_value. This lets you define custom logic for what gets returned when the mock is called.

Option 1: Use a Dictionary for Simple Mappings #

If your argument-return pairs are straightforward, a dictionary works perfectly. We'll wrap it in a lambda to fetch the right value:

@patch("my_module.check_item")
    def test_check_item_multiple_args(self, mock_check_item):
        # Define your argument-to-return mapping
        return_mapping = {
            "apple": "processed_apple",
            "orange": "processed_orange"
        }
        
        # Set side_effect to a lambda that looks up the input in the mapping
        mock_check_item.side_effect = lambda item: return_mapping[item]
        
        # Test apple case
        apple_result = my_module.process_fruit("apple")
        self.assertEqual(apple_result, "processed_apple")
        
        # Test orange case
        orange_result = my_module.process_fruit("orange")
        self.assertEqual(orange_result, "processed_orange")
        
        # Assert both calls happened in the right order
        mock_check_item.assert_has_calls([
            unittest.mock.call("apple"),
            unittest.mock.call("orange")
        ])
        # Optional: Assert total number of calls
        self.assertEqual(mock_check_item.call_count, 2)

Option 2: Use a Custom Function for Complex Logic #

If you need more control (like handling edge cases or multiple parameters), write a custom function for side_effect:

@patch("my_module.check_item")
    def test_check_item_complex_logic(self, mock_check_item):
        def custom_return(item):
            if item == "apple":
                return "processed_apple"
            elif item == "orange":
                return "processed_orange"
            else:
                return "unknown_item"
        
        mock_check_item.side_effect = custom_return
        
        # Test your cases
        self.assertEqual(my_module.process_fruit("apple"), "processed_apple")
        self.assertEqual(my_module.process_fruit("orange"), "processed_orange")
        self.assertEqual(my_module.process_fruit("banana"), "unknown_item")

Key Takeaways #

• Order matters: Always configure your mock first (set return_value or side_effect), then run the code that uses the mocked function, and finally assert the results and call history.
• side_effect is your friend: Use it for dynamic returns based on input, instead of static return_value.
• Assert multiple calls: Use assert_has_calls to verify the mock was called with the right arguments in the correct order.

内容的提问来源于stack exchange，提问作者Danny