| 54 | class PipelineInterfaceTest(unittest.TestCase): |
| 55 | |
| 56 | def prepare_dir(self, dirname, pipeline_name): |
| 57 | if not os.path.exists(dirname): |
| 58 | os.makedirs(dirname) |
| 59 | cfg_file = os.path.join(dirname, ModelFile.CONFIGURATION) |
| 60 | cfg = { |
| 61 | ConfigFields.framework: Frameworks.torch, |
| 62 | ConfigFields.task: Tasks.image_classification, |
| 63 | ConfigFields.pipeline: { |
| 64 | 'type': pipeline_name, |
| 65 | } |
| 66 | } |
| 67 | io.dump(cfg, cfg_file) |
| 68 | |
| 69 | def setUp(self) -> None: |
| 70 | self.prepare_dir('/tmp/custom_single_model', 'custom_single_model') |