diff --git a/model-archiver/model_archiver/model_packaging.py b/model-archiver/model_archiver/model_packaging.py index 345578dd34..52d2a9caf6 100644 --- a/model-archiver/model_archiver/model_packaging.py +++ b/model-archiver/model_archiver/model_packaging.py @@ -5,6 +5,7 @@ import logging import sys +import shutil from .arg_parser import ArgParser from .model_packaging_utils import ModelExportUtils from .model_archiver_error import ModelArchiverError @@ -36,7 +37,7 @@ def package_model(args, manifest): # Step 2 : Zip 'em all up ModelExportUtils.archive(export_file_path, model_name, model_path, manifest, args.archive_format) - + shutil.rmtree(model_path) logging.info("Successfully exported model %s to file %s", model_name, export_file_path) except ModelArchiverError as e: logging.error(e) diff --git a/model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py b/model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py index ccd847a488..6297318704 100644 --- a/model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py +++ b/model-archiver/model_archiver/tests/unit_tests/test_model_packaging.py @@ -3,6 +3,11 @@ from collections import namedtuple import pytest +import sys +from mock import MagicMock + +sys.modules['shutil'] = MagicMock() +sys.modules['shutil.rmtree'] = MagicMock() from model_archiver.manifest_components.manifest import RuntimeType from model_archiver.model_packaging import generate_model_archive, package_model