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

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 (имя = '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