如何用列表模拟/替换类

Question

基本上我想这样做:

class SeaLion(object):
    def __iter__(self):
        return self
    def next(self):
        ...
def train():
    sea_lion = SeaLion(...)
    for tooth in sea_lion:
        .. do something

然后在我的测试文件中我有:

class TestTrain(TestCase):
    @patch('module.SeaLion')
    def test_train(self, mock_SeaLion):
        mock_SeaLion.return_value = [1,2,3]
        ...

基本上，我希望train函数中的sea_lion是一个列表，而不是该类的实例。

上面的代码sea_lion是MagicMock(name='SeaLion()')

Answer 1

试试这个:

import unittest
from mock import patch
class SeaLion(object):
    def __iter__(self):
        return self
    def next(self):
        pass
def train():
    sea_lion = SeaLion()
    print(repr(sea_lion))
    for tooth in sea_lion:
        return True
    return False
class TestTrain(unittest.TestCase):
    @patch('__main__.SeaLion')
    def test_train(self, mock_SeaLion):
        mock_SeaLion.return_value = [1]
        self.assertTrue(train())
if __name__ == '__main__':
    unittest.main()

结果:

[1]
.
----------------------------------------------------------------------
Ran 1 test in 0.044s
OK