如何使用Mock框架更改模擬方法的輸出？

Question

我有類似下面的模塊：

'''example.py''' 
from some.module import Session 

def init_system(): 
    var_one = Session.query(SQLAModelA).all() 
    var_two = [(x.attr_B_a, x) for x in Session.query(SQLAModelB).all()] 

    for value in var_one: 
     something = value.attr_A_a + value.attr_A_c 
    # Will do something with var_two etc (which has different attributes) 
    something_else = [x.attr_B * 2 for x in var_two] 
    return result 

會議是指會談到一個數據庫的SQLAlchemy的會話處理程序。出於測試目的，我需要能夠模擬Session對象中的查詢方法（或者可能只是Session對象？），所以我可以測試init_system，而不實際與數據庫進行交互。我如何使用Mock framework來做到這一點？

真的讓我感到困惑的是，如果我嘲笑query()方法我該如何改變它的所有方法的輸出依賴於傳遞給query()的類？

我希望我能線沿線的做一些事情：

def test_init_numbers(): 
    import 
    my_mock_object = MagicMock() 
    with patch('some.module.Session.query', my_mock_object): 
     result = example.init_system() 
     assert result == ['expected', 'result'] 

不過，我不知道我需要做的my_mock_object讓它模擬query()正確。我想它需要實現返回一個迭代器的all()方法。而且，我知道我可以創造一些更模仿對象在返回的迭代中去，確保他們有值的相應屬性，如：

var_one_element_mock = MagicMock(spec=SQLAModelA) 
var_one_element_mock.attr_A_a = 12 

var_two_element_mock = MagicMock(spec=SQLAModelB) 
var_one_element_mock.attr_B = 100 

，當然，我可以創造不同的iterables那些可以被init_system函數使用的對象的版本，例如for value in var_one行。

但我不知道如何將所有模擬對象拼在一起，因此我可以剔除Session.query方法，同時考慮到將類傳遞給它，這會影響它返回的可迭代內容。

謝謝！

Answer 1

我不知道你是否仍然需要這方面的幫助，但是如果你想控制一個MagicMock對象被調用時會發生什麼，你可以傳遞一個函數作爲它的side_effect參數。 像這樣的東西應該工作。

def test_init_numbers(): 
    import 
    my_mock_object = MagicMock() 
    def return_iterator(): 
     ## create and return your iterator here 

    my_mock_object.all = MagicMock(side_effect=return_iterator) 
    with patch('some.module.Session.query', my_mock_object): 
     result = example.init_system() 
     assert result == ['expected', 'result']