16 lines
464 B
Python

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