from image_prediction.estimator.adapter.adapter import EstimatorAdapter class EstimatorMock: @staticmethod def predict(batch): return [None for _ in batch] class EstimatorAdapterMock(EstimatorAdapter): def __init__(self, estimator): super().__init__(estimator=estimator) def predict(self, batch): return self.estimator.predict(batch)