1616from mock import Mock , MagicMock , patch
1717
1818from sagemaker import chainer , estimator , model , mxnet , tensorflow , transformer , tuner , processing
19+ from sagemaker .network import NetworkConfig
1920from sagemaker .processing import ProcessingInput , ProcessingOutput
2021from sagemaker .workflow import airflow
2122from sagemaker .amazon import amazon_estimator
@@ -1598,6 +1599,13 @@ def test_deploy_config_from_amazon_alg_estimator(sagemaker_session):
15981599@patch ("sagemaker.utils.sagemaker_timestamp" , MagicMock (return_value = TIME_STAMP ))
15991600def test_processing_config (sagemaker_session ):
16001601
1602+ network_config = NetworkConfig (
1603+ encrypt_inter_container_traffic = False ,
1604+ enable_network_isolation = True ,
1605+ security_group_ids = ["sg1" ],
1606+ subnets = ["subnet1" ],
1607+ )
1608+
16011609 processor = processing .Processor (
16021610 role = "arn:aws:iam::0122345678910:role/SageMakerPowerUser" ,
16031611 image_uri = "{{ image_uri }}" ,
@@ -1612,6 +1620,7 @@ def test_processing_config(sagemaker_session):
16121620 sagemaker_session = sagemaker_session ,
16131621 tags = [{"{{ key }}" : "{{ value }}" }],
16141622 env = {"{{ key }}" : "{{ value }}" },
1623+ network_config = network_config ,
16151624 )
16161625
16171626 outputs = [
@@ -1699,5 +1708,10 @@ def test_processing_config(sagemaker_session):
16991708 "RoleArn" : "arn:aws:iam::0122345678910:role/SageMakerPowerUser" ,
17001709 "StoppingCondition" : {"MaxRuntimeInSeconds" : 3600 },
17011710 "Tags" : [{"{{ key }}" : "{{ value }}" }],
1711+ "NetworkConfig" : {
1712+ "EnableInterContainerTrafficEncryption" : False ,
1713+ "EnableNetworkIsolation" : True ,
1714+ "VpcConfig" : {"SecurityGroupIds" : ["sg1" ], "Subnets" : ["subnet1" ]},
1715+ },
17021716 }
17031717 assert config == expected_config
0 commit comments