from image_prediction.model_loader.loader import ModelLoader


class ModelLoaderMock(ModelLoader):

    model = None
    classes = None

    def load_model(self, identifier):
        assert self.model is not None, "Set the model to be returned first via monkeypatching"
        return self.model

    def load_classes(self, identifier):
        assert self.classes is not None, "Set the classes to be returned first via monkeypatching"
        return self.classes