Correct indentation and use placeholder for experiment_config in unit…

… test
aws · Oct 15, 2021 · 1422243 · 1422243
1 parent 3d7970a
commit 1422243
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 17 deletions.
diff --git a/tests/integ/test_sagemaker_steps.py b/tests/integ/test_sagemaker_steps.py
@@ -211,17 +211,17 @@ def test_transform_step_with_placeholder(trained_estimator, sfn_client, sfn_role
     })
 
     parameters = {
-            'BatchStrategy': execution_input['strategy'],
-            'TransformInput': {
-                'SplitType': execution_input['split_type'],
-            },
-            'TransformResources': {
-                'InstanceCount': execution_input['instance_count'],
-                'InstanceType': execution_input['instance_type'],
-            },
-            'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
-            'MaxPayloadInMB': execution_input['max_payload']
-        }
+        'BatchStrategy': execution_input['strategy'],
+        'TransformInput': {
+            'SplitType': execution_input['split_type'],
+        },
+        'TransformResources': {
+            'InstanceCount': execution_input['instance_count'],
+            'InstanceType': execution_input['instance_type'],
+        },
+        'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
+        'MaxPayloadInMB': execution_input['max_payload']
+    }
 
     # Build workflow definition
     transform_step = TransformStep(

diff --git a/tests/unit/test_sagemaker_steps.py b/tests/unit/test_sagemaker_steps.py
@@ -925,6 +925,7 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
         'tags': [{str: str}],
         'env': str,
         'volume_kms_key': str,
+        'experiment_config': str,
     })
 
     step_input = StepInput(schema={
@@ -944,10 +945,11 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
                 'InstanceType': step_input['instance_type'],
                 'VolumeKmsKeyId': execution_input['volume_kms_key']
             },
+            'ExperimentConfig': execution_input['experiment_config'],
             'Tags': execution_input['tags'],
             'Environment': execution_input['env'],
             'MaxConcurrentTransforms': execution_input['max_concurrent_transforms'],
-            'MaxPayloadInMB': execution_input['max_payload']
+            'MaxPayloadInMB': execution_input['max_payload'],
         }
 
     step = TransformStep('Inference',
@@ -997,11 +999,7 @@ def test_transform_step_creation_with_placeholder(pca_transformer):
             'InstanceType.$': "$['instance_type']",
             'VolumeKmsKeyId.$': "$$.Execution.Input['volume_kms_key']"
         },
-        'ExperimentConfig': {
-            'ExperimentName': 'pca_experiment',
-            'TrialName': 'pca_trial',
-            'TrialComponentDisplayName': 'Transform'
-        },
+        'ExperimentConfig.$': "$$.Execution.Input['experiment_config']",
         'DataProcessing': {
             'InputFilter.$': "$$.Execution.Input['input_filter']",
             'OutputFilter.$': "$$.Execution.Input['output_filter']",