你如何在for循环中模拟一个函数

时间:2018-02-01 07:58:11

标签: python-2.7 unit-testing mocking

我正在使用python的mock和补丁库。我有一个需要进行单元测试的功能。

我想进行单元测试func& mock / patch get_me_info函数迭代键A, B, C,每个键返回一个不同的数据帧输出。

因此get_me_info(pro_id='A')将数据框作为

返回
   src_id    date_info
0  A_src1    <date-string1>
1  A_src2    <date-string2>

因此get_me_info(pro_id='B')将数据框作为

返回
   src_id    date_info
0  B_src1    <date-string1>
1  B_src2    <date-string2>

需要测试下面的func():

def func():  
   result = []     
   for pro_id in [p.upper() for p in INFO]: # pro_id corresponds to here A/B/C keys in INFO

       df = cat_svc.get_me_info(pro_id=pro_id) # this call needs to be mocked

       for _, row in df.iterrows():
           # do some work that populates result & needs to be unit tested


   return result


INFO = {
        "A": {
            "s_key1": "a_val1",
            "s_key2": 11},
        "B": {
            "s_key1": "a_val1",
            "s_key2": 31},
        "C": {
            "s_key1": "b_val1",
            "s_key2": 23},
}

我应该如何通过get_me_info为测试func和此dict INFO返回不同的模拟数据帧?

1 个答案:

答案 0 :(得分:0)

我通过创建for循环的新函数并模拟该功能来解决这个问题。所以我现在测试了bar函数而不是func,并且能够模拟df = cat_svc.get_me_info(pro_id=pro_id)

# These functions are inside MyClass

    def func():  
       result = []     
       for pro_id in [p.upper() for p in INFO]: # pro_id corresponds to here A/B/C keys in INFO

           result.append(self.bar(pro_id))


       return result


    def bar(pro_id):
       res = []
       cat_svc.get_me_info(pro_id=pro_id) # this call needs to be mocked

           for _, row in df.iterrows():
               # do some work that populates res & needs to be unit tested
               #res.append(row)


        return res

这是测试:

   test_info = [
        {'key1': 'val1', 'key2': u'val2'},
     ]

    @mock.patch('a.b.get_me_info')
    @mock.patch('c.d.init')
    def test_bar(self, init_mock, me_mock_info):
        my_class = MyClass()
        me_mock_info.return_value = pd.DataFrame(test_info)
        result = my_class.bar(pro_id='test_id')
        self.assertIsNotNone(result)
        me_mock_info.assert_called_with(pro_id='test_id')