Как смоделировать / заменить класс списком

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]
        ...

По сути, я хочу, чтобы sea_lion в функции train был списком, а не экземпляром этого класса.

С указанным выше кодом 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