16
16
from mock import Mock , MagicMock , patch
17
17
18
18
from sagemaker import chainer , estimator , model , mxnet , tensorflow , transformer , tuner , processing
19
+ from sagemaker .network import NetworkConfig
19
20
from sagemaker .processing import ProcessingInput , ProcessingOutput
20
21
from sagemaker .workflow import airflow
21
22
from sagemaker .amazon import amazon_estimator
@@ -1598,6 +1599,13 @@ def test_deploy_config_from_amazon_alg_estimator(sagemaker_session):
1598
1599
@patch ("sagemaker.utils.sagemaker_timestamp" , MagicMock (return_value = TIME_STAMP ))
1599
1600
def test_processing_config (sagemaker_session ):
1600
1601
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
+
1601
1609
processor = processing .Processor (
1602
1610
role = "arn:aws:iam::0122345678910:role/SageMakerPowerUser" ,
1603
1611
image_uri = "{{ image_uri }}" ,
@@ -1612,6 +1620,7 @@ def test_processing_config(sagemaker_session):
1612
1620
sagemaker_session = sagemaker_session ,
1613
1621
tags = [{"{{ key }}" : "{{ value }}" }],
1614
1622
env = {"{{ key }}" : "{{ value }}" },
1623
+ network_config = network_config ,
1615
1624
)
1616
1625
1617
1626
outputs = [
@@ -1699,5 +1708,10 @@ def test_processing_config(sagemaker_session):
1699
1708
"RoleArn" : "arn:aws:iam::0122345678910:role/SageMakerPowerUser" ,
1700
1709
"StoppingCondition" : {"MaxRuntimeInSeconds" : 3600 },
1701
1710
"Tags" : [{"{{ key }}" : "{{ value }}" }],
1711
+ "NetworkConfig" : {
1712
+ "EnableInterContainerTrafficEncryption" : False ,
1713
+ "EnableNetworkIsolation" : True ,
1714
+ "VpcConfig" : {"SecurityGroupIds" : ["sg1" ], "Subnets" : ["subnet1" ]},
1715
+ },
1702
1716
}
1703
1717
assert config == expected_config
0 commit comments