如何在 python unittest 中永久模拟函数的返回值

vgb*_*ell 6 python patch mocking python-3.x python-unittest

我有一个函数

# foo.py
NUM_SHARDS = 10
def get_shard(shard_key: int) -> int
    return (shard_key % NUM_SHARDS) + 1
Run Code Online (Sandbox Code Playgroud)

我想模拟这个函数,以便每当调用这个函数时,它都会返回某个值。我试过像这样打补丁

# test.py
@mock.patch('foo.get_shard', return_value=11)
    def testSomething(self, patched_func):
        patched_func() # >> 11 ... OK so this is fine but
        get_shard(4) # >> 5 .... Original impl being executed
Run Code Online (Sandbox Code Playgroud)

我想深入修改这个函数以始终返回这个值,无论这个函数是否直接从单元测试调用。

作为上下文,我有许多生产分片 (10) 和一个测试分片 (11)。get_shard 函数被多次调用,该函数的生产版本仅支持 10 个分片。因此,我希望该函数在单元测试上下文中执行时始终返回 11(测试分片编号)。

blh*_*ing 8

mock.patch方法旨在暂时模拟给定的对象。如果你想永久改变函数的返回值,你可以简单地通过用指定的Mock对象分配给它来覆盖它return_value

import foo
from unittest.mock import Mock
foo.get_shard = Mock(return_value=11)
Run Code Online (Sandbox Code Playgroud)