def test_onnx_file_is_none(self, patches): patches.utils.find_unique.return_value = None ModelExportUtils.check_custom_model_types( model_path=self.model_path, model_name=None) patches.utils.find_unique.assert_called() patches.utils.convert_onnx_model.assert_not_called()
def test_onnx_file_is_not_none(self, patches): onnx_file = 'some-file.onnx' patches.utils.find_unique.return_value = onnx_file patches.utils.convert_onnx_model.return_value = ('sym', 'param') temp, exclude = ModelExportUtils.check_custom_model_types( self.model_path) patches.utils.convert_onnx_model.assert_called_once_with( self.model_path, onnx_file, None) assert len(temp) == 2 assert len(exclude) == 1 assert temp[0] == os.path.join(self.model_path, 'sym') assert temp[1] == os.path.join(self.model_path, 'param') assert exclude[0] == onnx_file