update tests to address breaking changes

Author: Malav Shastri

test/unit/test_handler_service.py

@@ -16,7 +16,7 @@
 import pytest
 import xgboost as xgb
 from mock import patch
-from sagemaker_containers.beta.framework import content_types, encoders, errors
+from sagemaker_training import content_types, encoders, errors
 
 from sagemaker_xgboost_container.handler_service import HandlerService
 

test/unit/test_serving.py

@@ -18,7 +18,7 @@
 import pytest
 import xgboost as xgb
 from mock import MagicMock, patch
-from sagemaker_containers.beta.framework import content_types, encoders, errors
+from sagemaker_training import content_types, encoders, errors
 
 from sagemaker_algorithm_toolkit.exceptions import UserError
 from sagemaker_xgboost_container import serving

test/utils/local_mode.py

@@ -27,7 +27,7 @@
 import yaml
 from botocore.exceptions import ClientError
 from sagemaker import fw_utils, utils
-from sagemaker_containers.beta.framework import content_types
+from sagemaker_training import content_types
 
 CYAN_COLOR = "\033[36m"
 END_COLOR = "\033[0m"