Merge pull request aws#23 from verdimrc/fp-mxnet-uni-test

added a unit test for mxnet processor

Author: verdimrc

tests/unit/test_processing.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 from __future__ import absolute_import
+from platform import python_version
 
 import pytest
 from mock import Mock, patch, MagicMock
@@ -33,6 +34,7 @@
 from sagemaker.pytorch.processing import PyTorchProcessor
 from sagemaker.utils import get_config_value
 from sagemaker.xgboost.processing import XGBoostEstimator
+from sagemaker.mxnet.processing import MXNetProcessor
 from sagemaker.network import NetworkConfig
 from sagemaker.processing import FeatureStoreOutput
 from sagemaker.fw_utils import UploadedCode
@@ -390,6 +392,44 @@ def test_xgboost_with_required_parameters(
 
     sagemaker_session.process.assert_called_with(**expected_args)
 
+@patch("sagemaker.utils._botocore_resolver")
+@patch("os.path.exists", return_value=True)
+@patch("os.path.isfile", return_value=True)
+def test_mxnet_with_required_parameters(
+    exists_mock, isfile_mock, botocore_resolver, sagemaker_session, mxnet_training_version, mxnet_training_py_version
+):
+    botocore_resolver.return_value.construct_endpoint.return_value = {"hostname": ECR_HOSTNAME}
+
+    processor = MXNetProcessor(
+        role=ROLE,
+        instance_type="ml.m4.xlarge",
+        framework_version=mxnet_training_version,
+        py_version=mxnet_training_py_version,
+        instance_count=1,
+        sagemaker_session=sagemaker_session,
+    )
+
+    processor.run(code="/local/path/to/processing_code.py")
+
+    expected_args = _get_expected_args_modular_code(processor._current_job_name)
+
+    if (mxnet_training_py_version == 'py3') & (mxnet_training_version == "1.4"): #probably there is a better way to handle this
+        mxnet_image_uri = (
+            "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-training:{}-cpu-{}"
+        ).format(mxnet_training_version, mxnet_training_py_version)
+    elif version.parse(mxnet_training_version) > version.parse("1.4.1" if mxnet_training_py_version == 'py2' else "1.4"):
+        mxnet_image_uri = (
+            "763104351884.dkr.ecr.us-west-2.amazonaws.com/mxnet-training:{}-cpu-{}"
+        ).format(mxnet_training_version, mxnet_training_py_version)
+    else:
+        mxnet_image_uri = (
+            "520713654638.dkr.ecr.us-west-2.amazonaws.com/sagemaker-mxnet:{}-cpu-{}"
+        ).format(mxnet_training_version, mxnet_training_py_version)
+    
+    expected_args["app_specification"]["ImageUri"] = mxnet_image_uri
+
+    sagemaker_session.process.assert_called_with(**expected_args)
+
 
 @patch("os.path.exists", return_value=False)
 def test_script_processor_errors_with_nonexistent_local_code(exists_mock, sagemaker_session):