如何在Python中对ProcessPoolExecutor启动的子进程内的elasticsearch.Elasticsearch进行Mock？

我完全懂你碰到的这个坑——主进程里用monkeypatch搞的Mock明明好好的，一到ProcessPoolExecutor启动的子进程里就失效了，这其实是Python多进程的特性在搞鬼。

为什么会出现这个问题？ #

ProcessPoolExecutor创建子进程的时候，不管是用Unix系统默认的fork模式，还是Windows/显式指定的spawn模式，子进程都会独立启动一个Python解释器，并且重新导入所有需要的模块。主进程里通过monkeypatch修改的elasticsearch.Elasticsearch属性，并不会自动传递给子进程：

• 直接调用tt.just_es()时，代码在主进程执行，Mock自然生效；
• 但用executor.submit启动子进程后，子进程会重新加载elasticsearch模块，拿到的是原始的Elasticsearch类，所以Mock就没效果了。

解决方案：给子进程加初始化钩子 #

咱们可以利用ProcessPoolExecutor自带的initializer和initargs参数，让每个子进程启动后自动执行Mock逻辑，把Elasticsearch替换成我们的假客户端。下面分两种情况给出实现方式：

方式一：调整主代码（更灵活）

先修改main.py里的process方法，允许传入初始化函数，这样测试时不用硬编码Mock逻辑：

import concurrent.futures
import elasticsearch

class ThreadTest:
    def __init__(self):
        pass

    def process(self, initializer=None, initargs=None):
        futures = []
        # 把初始化函数传给ProcessPoolExecutor
        with concurrent.futures.ProcessPoolExecutor(
            max_workers=5,
            initializer=initializer,
            initargs=initargs or ()
        ) as executor:
            future = executor.submit(self.just_es)
            futures.append(future)

        concurrent.futures.wait(fs=futures)
        for future in futures:
            result = future.result()

    def just_es(self):
        print(elasticsearch.Elasticsearch)
        return "hey"

然后修改test.py，定义子进程的初始化函数，并在测试时传入：

import elasticsearch
import pytest
from main import ThreadTest


class MockElasticsearchClient:
    def __init__(self):
        pass

    @property
    def meta(self):
        return


def fake_es_client_6_for_run_query(*args, **kwargs):
    return MockElasticsearchClient()


@pytest.fixture(autouse=True)
def mock_main_process_es(monkeypatch):
    # 给主进程打Mock
    monkeypatch.setattr(
        elasticsearch,
        "Elasticsearch",
        fake_es_client_6_for_run_query,
    )


def init_subprocess_mock():
    # 子进程启动时执行这个函数，给子进程打Mock
    import elasticsearch
    elasticsearch.Elasticsearch = fake_es_client_6_for_run_query


def test_thread(mock_main_process_es):
    tt = ThreadTest()
    tt.just_es()  # 主进程Mock生效
    # 传入初始化函数，让子进程自动应用Mock
    tt.process(initializer=init_subprocess_mock)

方式二：不修改主代码（更侵入性小）

如果你不想改动业务代码，可以用unittest.mock.patch动态替换ProcessPoolExecutor的初始化逻辑，偷偷加上initializer参数：

import elasticsearch
import pytest
from unittest.mock import patch
from main import ThreadTest
import concurrent.futures


class MockElasticsearchClient:
    def __init__(self):
        pass

    @property
    def meta(self):
        return


def fake_es_client_6_for_run_query(*args, **kwargs):
    return MockElasticsearchClient()


@pytest.fixture(autouse=True)
def mock_main_process_es(monkeypatch):
    monkeypatch.setattr(
        elasticsearch,
        "Elasticsearch",
        fake_es_client_6_for_run_query,
    )


def init_subprocess_mock():
    import elasticsearch
    elasticsearch.Elasticsearch = fake_es_client_6_for_run_query


def test_thread(mock_main_process_es):
    tt = ThreadTest()
    tt.just_es()

    # 用patch替换ProcessPoolExecutor，自动添加初始化逻辑
    with patch('concurrent.futures.ProcessPoolExecutor') as mock_executor:
        def executor_factory(max_workers=5):
            return concurrent.futures.ProcessPoolExecutor(
                max_workers=max_workers,
                initializer=init_subprocess_mock
            )
        mock_executor.side_effect = executor_factory
        tt.process()

验证效果 #

运行测试后，你会看到主进程和子进程的输出都是我们的Mock函数，说明Mock在子进程里也生效了：

<function fake_es_client_6_for_run_query at 0x102703c40>
<function fake_es_client_6_for_run_query at 0x102703c40>

备注：内容来源于stack exchange，提问作者Shashwat