update tests to address breaking changes

aws · Aug 13, 2024 · aacfaf8 · aacfaf8
1 parent 6881b23
commit aacfaf8
Show file tree

Hide file tree

Showing 3 changed files with 3 additions and 3 deletions.
diff --git a/test/unit/test_handler_service.py b/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
 

diff --git a/test/unit/test_serving.py b/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

diff --git a/test/utils/local_mode.py b/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"