Python测试中如何Mock主脚本依赖的SQLES类：解决自定义Mock类未生效问题

I get exactly what you're trying to do—you want to test main.py in isolation, using your custom MockSQLES instead of the real SQLES class, without touching a single line of main.py. The issue with your current code is that unittest.mock.patch defaults to replacing the target with a MagicMock instead of your custom class. Let's fix that.

Why Your Current Code Isn't Working #

When you write with unittest.mock.patch('main.SQLES') as MockSQLES:, you're telling Python to replace main.SQLES with a MagicMock instance. So when main.py runs sql = SQLES(), it's creating a MagicMock object, not an instance of your MockSQLES class. That's why your custom methods never get called.

The Fix: Patch with Your Custom Mock Class #

You need to explicitly tell patch to use your MockSQLES class as the replacement. There are two straightforward ways to do this:

Method 1: Use the new Parameter

This directly replaces the target class with your mock class. Here's how to adjust your test code:

import unittest.mock
import main

class MockSQLES:
    def __init__(self):
        # Initialize attributes to match the real class structure
        self.exchange_rates = {}
        self.firms_names_other = {}

    def init(self, logger, cmd_args_record, cmd_args_playback):
        return

    def get_exchange_rates(self):
        self.exchange_rates['GBP'] = [
            ['2021-09-13', 0.722298],
            ['2021-08-20', 0.734203]
        ]
        return

    def get_firms_names_other(self):
        # Fixed the undefined `firm_id` issue (adjust to match your actual use case)
        firm_id = "sample_firm_id"
        self.firms_names_other[firm_id] = [
            [131, 'Acme Industries'],
            [132, 'Acme Research']
        ]
        return

def test_main_with_mock():
    # Replace main.SQLES with our custom MockSQLES class
    with unittest.mock.patch('main.SQLES', new=MockSQLES):
        # Now when main.py runs sql = SQLES(), it creates a MockSQLES instance
        main.main()
        
        # Optional: Add assertions to verify main.py used the mock data correctly
        # For example, if main.py exposes the sql instance:
        # assert main.sql.exchange_rates['GBP'] == [['2021-09-13', 0.722298], ['2021-08-20', 0.734203]]

if __name__ == '__main__':
    test_main_with_mock()

Method 2: Use side_effect

If you need more control over instance creation (e.g., custom initialization logic per instance), use side_effect. When set to a class, it returns a new instance of your mock every time the target is instantiated:

with unittest.mock.patch('main.SQLES', side_effect=MockSQLES):
    main.main()

Key Notes to Avoid Headaches #

• Patch the Right Path: Always patch main.SQLES (the reference inside main.py), not sqles.SQLES (the original class). patch replaces the reference that main.py actually uses, so targeting the wrong path won't work.
• Match the Original Interface: Ensure your MockSQLES has the same method names, parameters, and attribute structure as the real SQLES. This keeps main.py running without errors.
• Automatic Cleanup: Unlike manual monkey-patching, unittest.mock.patch restores the original SQLES class once the with block exits, keeping your tests isolated.

Would Monkey-Patching Work? #

Yes, you could manually replace the class like this:

original_sqles = main.SQLES
main.SQLES = MockSQLES
try:
    main.main()
finally:
    main.SQLES = original_sqles

But unittest.mock.patch is better—it handles cleanup automatically and integrates smoothly with test frameworks like unittest or pytest.

内容的提问来源于stack exchange，提问作者John R Ramsden