class EstimatorMock:
    def __init__(self):
        self.__output_batch = None

    @property
    def output_batch(self):
        return self.__output_batch

    @output_batch.setter
    def output_batch(self, output_batch):
        self.__output_batch = output_batch

    def predict(self, batch):
        return self.__output_batch