如何在 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()。

替代解决方案

在这种情况下，如果您想测试random 函数的行为，最好使用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

Answer 2

我不知道测试中的 mockobj 是什么，但你可以做的是。

    @mock.patch('random.choice')
    def test1(self, choice_mock):
        choice_mock.return_value = 1000
        self.assertEqual(a(), 1000)

Answer 3

我想用完整的脚本改进@Alex 响应，以便更好地理解和适应其他情况。

import random
from unittest import TestCase, mock

letters = ['A', 'B', 'C', 'D']

def get_random_words(): # Simple function using choice
  l = []
  for _ in range(3):
      l.append(random.choice(letters))
    
  return "".join(l)

class TestRandom(TestCase):

   @mock.patch('random.choice') # *(1)
   def test_get_random_words(self, mock_choice):
    
      mock_choice.side_effect = ['A','b','D','Z'] # *(2)
      result = get_random_words()
    
      self.assertEqual(result, 'AbD', 'Does not generate correct string')

注意事项

*(1) 对于本例，函数在同一个文件中，但如果它在另一个文件中，则必须更改补丁的路径 例如：@mock.patch('your_package.your_file.your_function.random.choice')

*(2) 对于这种情况，get_random_words 函数调用random.choice 3 次。这就是为什么您必须在 mock_choice.side_effect 中放入相同或更多的项目。这是因为如果它的项目较少，它将引发StopIteration 错误。