python - 如何在 Python 的单元测试框架中模拟返回 self 的方法

Question

我正在使用一个类，该类具有一个方法shuffle，该方法返回调用它的实例的打乱版本。这是:

shuffled_object = unshuffled_object.shuffle(buffer_size)

我想模拟这个方法，以便在调用它时，它只返回自身，而不进行任何改组。以下是这种情况的简化:

# my_test.py
class Test():

    def shuffle(self, buffer_size):
        return self

# test_mock
import unittest
import unittest.mock as mk

import my_test

def mock_test(self, buffer_size):
    return self

class TestMock(unittest.TestCase):

    def test_mock(self):
        with mk.patch('my_test.Test.shuffle') as shuffle:
            shuffle.side_effect = mock_test
            shuffled_test = my_test.Test().shuffle(5)

但是，当我尝试此操作时，出现以下错误:

TypeError: mock_test() missing 1 required positional argument: 'buffer_size'

仅使用参数 5 调用该方法，调用实例不会将其自身作为 self 参数传递给该方法。是否可以使用 unittest.mock 模块实现这样的行为？

<小时/>编辑:

真正的代码是这样的:

# input.py
def create_dataset():
    ...
    raw_dataset = tf.data.Dataset.from_generator(data_generator, output_types, output_shapes)
    shuffled_dataset = raw_dataset.shuffle(buffer_size)
    dataset = shuffled_dataset.map(_load_example)
    ...
    return dataset

# test.py
def shuffle(self, buffer_size):
    return self

with mk.patch(input.tf.data.Dataset.shuffle) as shuffle_mock:
    shuffle_mock.side_effect = shuffle
    dataset = input.create_dataset()

这里的大问题是我只想模拟 shuffle 方法，因为我不希望它在测试时是随机的，但我想保留其余的原始方法，以便我的代码可以继续工作。棘手的部分是 shuffle 不仅对调用它的实例进行洗牌，而且还返回洗牌后的实例，因此我想在测试时返回数据集的未洗牌版本。

另一方面，让模拟继承tf.data.Dataset并不是那么简单，因为据我了解，Dataset似乎是一个具有抽象方法的抽象类，我想从初始化器 from_generator 创建的 Dataset 的任何子类型中抽象出来。

<小时/>编辑2:

我通过修补该方法进一步进行了如下操作:

def shuffle(cls, buffer_size, seed=None, reshuffle_each_iteration=None):
            def _load_example(example):
                return example
            return cls.map(cls, _load_example)

from data_input.kitti.kitti_input import tf as tf_mock

with mk.patch.object(tf_mock.data.Dataset, 'shuffle', classmethod(shuffle)):
    dataset = create_dataset()

现在，实例 raw_dataset 似乎将自身作为 shuffle 的 self 参数传递，但无论出于何种原因，代码仍然崩溃以下错误:

AttributeError: 'property' object has no attribute '_flat_types'

所以我假设这个self在某种程度上不完全是调用实例，它在某种程度上是内部不同的。

Answer 1

为什么没有 self 参数

声明类时，您定义的函数将在实例中绑定(bind)为方法。这是一个例子:

>>> def function():
...     pass
... 
>>> type(function)
<class 'function'>
>>> class A:
...     def b(self):
...         print(self)
>>> type(A.b)
<class 'function'>
>>> a = A()
>>> type(a.b)
<class 'method'>
# So you have the same behavior between the two following calls
>>> A.b(a)
<__main__.A object at 0x7f734511afd0>
>>> a.b()
<__main__.A object at 0x7f734511afd0>

解决方案

我可以提出一些解决方案，但并非所有解决方案都引人注目，具体取决于您的使用和需求。

模拟类

您可以模拟整个类来覆盖函数定义。如前所述，这考虑到您不使用类的抽象。

import unittest
import unittest.mock as mk

import my_test
import another

class TestMocked(my_test.Test):
    def shuffle(self, buffer_size):
        return self

@mk.patch("my_test.Test", TestMocked)
# Uncomment to mock the other file behavior
# @mk.patch("another.Test", TestMocked)
def test_mock():
    test_class = my_test.Test()
    shuffled_test = test_class.shuffle(2)
    print(my_test.Test.shuffle)
    # This is another file using your class,
    # You will have to mock it too in order to see the mocked behavior
    print(another.Test.shuffle) 
    assert shuffled_test == test_class

将输出:

>>> from test_mock import test_mock
>>> test_mock()
<function TestMocked.shuffle at 0x7ff1f03f0ae8>
<function Test.shuffle at 0x7ff1f03f09d8>

直接调用函数

我不喜欢这个，因为它会让你更改测试代码。您可以将调用从 instance.method() 转换为 class.method(instance)。这将按预期将参数发送到您的模拟函数。

# my_input.py
import tensorflow as tf


def data_generator():
    for i in itertools.count(1):
        yield (i, [1] * i)


def create_dataset():
    _load_example = lambda x, y: x+y
    buffer_size = 3
    output_types = (tf.int64, tf.int64)
    output_shapes = (tf.TensorShape([]), tf.TensorShape([None]))
    raw_dataset = tf.data.Dataset.from_generator(data_generator, output_types, output_shapes)

    shuffled_dataset = tf.data.Dataset.shuffle(raw_dataset, buffer_size)

    assert raw_dataset == shuffled_dataset
    assert raw_dataset is shuffled_dataset

    dataset = shuffled_dataset.map(_load_example)
    return dataset


# test_mock.py
import unittest.mock as mk
import my_input


def shuffle(self, buffer_size): 
    print("Shuffle! {}, {}".format(self, buffer_size))
    return self


with mk.patch('my_input.tf.data.Dataset.shuffle') as shuffle_mock:
    shuffle_mock.side_effect = shuffle
    dataset = my_input.create_dataset()

运行时，您将得到以下输出:

$ python test_mock.py
Shuffle! (<DatasetV1Adapter shapes: ((), (?,)), types: (tf.int64, tf.int64)>, 3)

将方法使用包装在函数中

这几乎与之前的答案相同，但是您可以将其包装如下，而不是从类中调用方法:

# my_input.py
import tensorflow as tf


def data_generator():
    for i in itertools.count(1):
        yield (i, [1] * i)


def shuffle(instance, buffer_size):
    return instance.shuffle(buffer_size)


def create_dataset():
    _load_example = lambda x, y: x+y
    buffer_size = 3
    output_types = (tf.int64, tf.int64)
    output_shapes = (tf.TensorShape([]), tf.TensorShape([None]))
    raw_dataset = tf.data.Dataset.from_generator(data_generator, output_types, output_shapes)

    shuffled_dataset = tf.data.Dataset.shuffle(raw_dataset, buffer_size)

    assert raw_dataset == shuffled_dataset
    assert raw_dataset is shuffled_dataset

    dataset = shuffled_dataset.map(_load_example)
    return dataset



# test_mock.py
import unittest.mock as mk
import my_input


def shuffle(self, buffer_size): 
    print("Shuffle! {}, {}".format(self, buffer_size))
    return self


with mk.patch('my_input.shuffle') as shuffle_mock:
    shuffle_mock.side_effect = shuffle
    dataset = my_input.create_dataset()