Skip to content

Commit c1fca41

Browse files
balajitummalaEthanShouhanCheng
authored andcommitted
fix: fix feature store ingestion via data wrangler test (aws#2687)
1 parent 2b0f764 commit c1fca41

File tree

1 file changed

+10
-1
lines changed

1 file changed

+10
-1
lines changed

tests/integ/test_workflow.py

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -2256,6 +2256,7 @@ def test_one_step_ingestion_pipeline(
22562256
input_data_uri = os.path.join(
22572257
"s3://", sagemaker_session.default_bucket(), "py-sdk-ingestion-test-input/features.csv"
22582258
)
2259+
22592260
with open(input_file_path, "r") as data:
22602261
body = data.read()
22612262
S3Uploader.upload_string_as_file_body(
@@ -2292,6 +2293,10 @@ def test_one_step_ingestion_pipeline(
22922293
)
22932294
]
22942295

2296+
output_content_type = "CSV"
2297+
output_config = {output_name: {"content_type": output_content_type}}
2298+
job_argument = [f"--output-config '{json.dumps(output_config)}'"]
2299+
22952300
temp_flow_path = "./ingestion.flow"
22962301
with cleanup_feature_group(feature_group):
22972302
json.dump(ingestion_only_flow, open(temp_flow_path, "w"))
@@ -2306,7 +2311,11 @@ def test_one_step_ingestion_pipeline(
23062311
)
23072312

23082313
data_wrangler_step = ProcessingStep(
2309-
name="ingestion-step", processor=data_wrangler_processor, inputs=inputs, outputs=outputs
2314+
name="ingestion-step",
2315+
processor=data_wrangler_processor,
2316+
inputs=inputs,
2317+
outputs=outputs,
2318+
job_arguments=job_argument,
23102319
)
23112320

23122321
pipeline = Pipeline(

0 commit comments

Comments
 (0)