Skip to content

Commit 1422243

Browse files
committed
Correct indentation and use placeholder for experiment_config in unit test
1 parent 3d7970a commit 1422243

File tree

2 files changed

+15
-17
lines changed

2 files changed

+15
-17
lines changed

Diff for: tests/integ/test_sagemaker_steps.py

+11-11
Original file line numberDiff line numberDiff line change
@@ -211,17 +211,17 @@ def test_transform_step_with_placeholder(trained_estimator, sfn_client, sfn_role
211211
})
212212

213213
parameters = {
214-
'BatchStrategy': execution_input['strategy'],
215-
'TransformInput': {
216-
'SplitType': execution_input['split_type'],
217-
},
218-
'TransformResources': {
219-
'InstanceCount': execution_input['instance_count'],
220-
'InstanceType': execution_input['instance_type'],
221-
},
222-
'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
223-
'MaxPayloadInMB': execution_input['max_payload']
224-
}
214+
'BatchStrategy': execution_input['strategy'],
215+
'TransformInput': {
216+
'SplitType': execution_input['split_type'],
217+
},
218+
'TransformResources': {
219+
'InstanceCount': execution_input['instance_count'],
220+
'InstanceType': execution_input['instance_type'],
221+
},
222+
'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
223+
'MaxPayloadInMB': execution_input['max_payload']
224+
}
225225

226226
# Build workflow definition
227227
transform_step = TransformStep(

Diff for: tests/unit/test_sagemaker_steps.py

+4-6
Original file line numberDiff line numberDiff line change
@@ -925,6 +925,7 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
925925
'tags': [{str: str}],
926926
'env': str,
927927
'volume_kms_key': str,
928+
'experiment_config': str,
928929
})
929930

930931
step_input = StepInput(schema={
@@ -944,10 +945,11 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
944945
'InstanceType': step_input['instance_type'],
945946
'VolumeKmsKeyId': execution_input['volume_kms_key']
946947
},
948+
'ExperimentConfig': execution_input['experiment_config'],
947949
'Tags': execution_input['tags'],
948950
'Environment': execution_input['env'],
949951
'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
950-
'MaxPayloadInMB': execution_input['max_payload']
952+
'MaxPayloadInMB': execution_input['max_payload'],
951953
}
952954

953955
step = TransformStep('Inference',
@@ -997,11 +999,7 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
997999
'InstanceType.$': "$['instance_type']",
9981000
'VolumeKmsKeyId.$': "$$.Execution.Input['volume_kms_key']"
9991001
},
1000-
'ExperimentConfig': {
1001-
'ExperimentName': 'pca_experiment',
1002-
'TrialName': 'pca_trial',
1003-
'TrialComponentDisplayName': 'Transform'
1004-
},
1002+
'ExperimentConfig.$': "$$.Execution.Input['experiment_config']",
10051003
'DataProcessing': {
10061004
'InputFilter.$': "$$.Execution.Input['input_filter']",
10071005
'OutputFilter.$': "$$.Execution.Input['output_filter']",

0 commit comments

Comments
 (0)