python - 如何在 python 中模拟 random.choice？

Question

我希望 choice 在我的单元测试中每次都返回相同的值 1000。以下代码不起作用。

import unittest
from random import choice

from mock import mock

def a():
    return choice([1, 2, 3])

class mockobj(object):
    @classmethod
    def choice(cls, li):
        return 1000

class testMock(unittest.TestCase):

    def test1(self):
        with mock.patch('random.choice', mockobj.choice):
            self.assertEqual(a(), 1000)

错误信息如下:

Failure
Traceback (most recent call last):
  File "test.py", line 15, in test1
    self.assertEqual(a(), 1000)
AssertionError: 3 != 1000

我应该如何修改它才能使其工作？我用的是python2.7

Answer 1

这里的问题是 a() 使用的是 未修补 版本的 random.choice。

比较函数 a 和 b:

import random
from random import choice

def a():
    return choice([1, 2, 3])

def b():
    return random.choice([1, 2, 3])

def choice1000(values):
    return 1000

import unittest.mock as mock

with mock.patch('random.choice', choice1000):
    print('a', a())
    print('b', b())

它打印出例如:

a 3
b 1000

为什么？

这一行是问题所在:

from random import choice

它导入了 random 并将 random.choice 存储到一个名为 choice 的新变量中。

后来，mock.patch 修补了原来的random.choice，但没有修补本地的choice。

我可以给本地打补丁吗？是的:

with mock.patch('__main__.choice', choice1000):
    print('a', a())
    print('b', b())

现在它打印出例如

a 1000
b 1

(我使用了 '__main__' 因为我将这段代码放入了 top-level file - 在你的情况下它可能是其他东西)

那怎么办？

要么修补一切，要么采取不同的方法。例如，修补 a() 而不是 choice()。

替代方案

在这种情况下，如果您想测试随机 函数的行为，最好使用seed。

def a():
    return random.choice([1, 2, 3, 1000])

def test1(self):
    random.seed(0)
    self.assertEqual(a(), 1000)

您无法事先知道某个种子会生成哪些随机值，但您可以确定它们始终相同。这正是您在测试中需要的。

在上面的最后一个例子中，我在 random.seed(0) 之后测试了 a() 一次，它返回了 1000，所以我可以肯定它会这样做每次:

>>> import random
>>> random.seed(0)
>>> print (random.choice([1, 2, 3, 1000]))
1000
>>> random.seed(0)
>>> print (random.choice([1, 2, 3, 1000]))
1000
>>> random.seed(0)
>>> print (random.choice([1, 2, 3, 1000]))
1000
>>> random.seed(0)
>>> print (random.choice([1, 2, 3, 1000]))
1000