Skip to content

Commit cd1fc9e

Browse files
committed
Removed integ test for retry and catch in constructor and updated existing tests
1 parent f3a0cc1 commit cd1fc9e

File tree

1 file changed

+44
-122
lines changed

1 file changed

+44
-122
lines changed

tests/integ/test_state_machine_definition.py

Lines changed: 44 additions & 122 deletions
Original file line numberDiff line numberDiff line change
@@ -422,51 +422,29 @@ def test_task_state_machine_creation(sfn_client, sfn_role_arn, training_job_para
422422

423423
def test_catch_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters):
424424
catch_state_name = "TaskWithCatchState"
425-
custom_error = "CustomError"
426425
task_failed_error = "States.TaskFailed"
427-
all_fail_error = "States.ALL"
428-
custom_error_state_name = "Custom Error End"
429-
task_failed_state_name = "Task Failed End"
430-
all_error_state_name = "Catch All End"
426+
timeout_error = "States.Timeout"
427+
task_failed_state_name = "Catch Task Failed End"
428+
timeout_state_name = "Catch Timeout End"
431429
catch_state_result = "Catch Result"
432430
task_resource = f"arn:{get_aws_partition()}:states:::sagemaker:createTrainingJob.sync"
433431

434-
# change the parameters to cause task state to fail
432+
# Provide invalid TrainingImage to cause States.TaskFailed error
435433
training_job_parameters["AlgorithmSpecification"]["TrainingImage"] = "not_an_image"
436434

437-
asl_state_machine_definition = {
438-
"StartAt": catch_state_name,
439-
"States": {
440-
catch_state_name: {
441-
"Resource": task_resource,
442-
"Parameters": training_job_parameters,
443-
"Type": "Task",
444-
"End": True,
445-
"Catch": [
446-
{
447-
"ErrorEquals": [
448-
all_fail_error
449-
],
450-
"Next": all_error_state_name
451-
}
452-
]
453-
},
454-
all_error_state_name: {
455-
"Type": "Pass",
456-
"Result": catch_state_result,
457-
"End": True
458-
}
459-
}
460-
}
461435
task = steps.Task(
462436
catch_state_name,
463437
parameters=training_job_parameters,
464-
resource=task_resource
438+
resource=task_resource,
439+
catch=steps.Catch(
440+
error_equals=[timeout_error],
441+
next_step=steps.Pass(timeout_state_name, result=catch_state_result)
442+
)
465443
)
466444
task.add_catch(
467445
steps.Catch(
468-
error_equals=[all_fail_error],
469-
next_step=steps.Pass(all_error_state_name, result=catch_state_result)
446+
error_equals=[task_failed_error],
447+
next_step=steps.Pass(task_failed_state_name, result=catch_state_result)
470448
)
471449
)
472450

@@ -476,19 +454,6 @@ def test_catch_state_machine_creation(sfn_client, sfn_role_arn, training_job_par
476454
role=sfn_role_arn
477455
)
478456

479-
workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, catch_state_result)
480-
481-
482-
def test_state_machine_creation_with_catch_in_constructor(sfn_client, sfn_role_arn, training_job_parameters):
483-
catch_state_name = "TaskWithCatchState"
484-
all_fail_error = "States.ALL"
485-
all_error_state_name = "Catch All End"
486-
catch_state_result = "Catch Result"
487-
task_resource = f"arn:{get_aws_partition()}:states:::sagemaker:createTrainingJob.sync"
488-
489-
# change the parameters to cause task state to fail
490-
training_job_parameters["AlgorithmSpecification"]["TrainingImage"] = "not_an_image"
491-
492457
asl_state_machine_definition = {
493458
"StartAt": catch_state_name,
494459
"States": {
@@ -500,78 +465,61 @@ def test_state_machine_creation_with_catch_in_constructor(sfn_client, sfn_role_a
500465
"Catch": [
501466
{
502467
"ErrorEquals": [
503-
all_fail_error
468+
timeout_error
469+
],
470+
"Next": timeout_state_name
471+
},
472+
{
473+
"ErrorEquals": [
474+
task_failed_error
504475
],
505-
"Next": all_error_state_name
476+
"Next": task_failed_state_name
506477
}
507478
]
508479
},
509-
all_error_state_name: {
480+
task_failed_state_name: {
510481
"Type": "Pass",
511482
"Result": catch_state_result,
512483
"End": True
513-
}
484+
},
485+
timeout_state_name: {
486+
"Type": "Pass",
487+
"Result": catch_state_result,
488+
"End": True
489+
},
514490
}
515491
}
516-
task = steps.Task(
517-
catch_state_name,
518-
parameters=training_job_parameters,
519-
resource=task_resource,
520-
catch=steps.Catch(
521-
error_equals=[all_fail_error],
522-
next_step=steps.Pass(all_error_state_name, result=catch_state_result)
523-
)
524-
)
525-
526-
workflow = Workflow(
527-
unique_name_from_base('Test_Catch_In_Constructor_Workflow'),
528-
definition=task,
529-
role=sfn_role_arn
530-
)
531492

532493
workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, catch_state_result)
533494

534495

535496
def test_retry_state_machine_creation(sfn_client, sfn_role_arn, training_job_parameters):
536497
retry_state_name = "RetryStateName"
537-
all_fail_error = "Starts.ALL"
498+
task_failed_error = "States.TaskFailed"
499+
timeout_error = "States.Timeout"
538500
interval_seconds = 1
539501
max_attempts = 2
540502
backoff_rate = 2
541503
task_resource = f"arn:{get_aws_partition()}:states:::sagemaker:createTrainingJob.sync"
542504

543-
# change the parameters to cause task state to fail
505+
# Provide invalid TrainingImage to cause States.TaskFailed error
544506
training_job_parameters["AlgorithmSpecification"]["TrainingImage"] = "not_an_image"
545507

546-
asl_state_machine_definition = {
547-
"StartAt": retry_state_name,
548-
"States": {
549-
retry_state_name: {
550-
"Resource": task_resource,
551-
"Parameters": training_job_parameters,
552-
"Type": "Task",
553-
"End": True,
554-
"Retry": [
555-
{
556-
"ErrorEquals": [all_fail_error],
557-
"IntervalSeconds": interval_seconds,
558-
"MaxAttempts": max_attempts,
559-
"BackoffRate": backoff_rate
560-
}
561-
]
562-
}
563-
}
564-
}
565-
566508
task = steps.Task(
567509
retry_state_name,
568510
parameters=training_job_parameters,
569-
resource=task_resource
511+
resource=task_resource,
512+
retry=steps.Retry(
513+
error_equals=[timeout_error],
514+
interval_seconds=interval_seconds,
515+
max_attempts=max_attempts,
516+
backoff_rate=backoff_rate
517+
)
570518
)
571519

572520
task.add_retry(
573521
steps.Retry(
574-
error_equals=[all_fail_error],
522+
error_equals=[task_failed_error],
575523
interval_seconds=interval_seconds,
576524
max_attempts=max_attempts,
577525
backoff_rate=backoff_rate
@@ -584,20 +532,6 @@ def test_retry_state_machine_creation(sfn_client, sfn_role_arn, training_job_par
584532
role=sfn_role_arn
585533
)
586534

587-
workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, None)
588-
589-
590-
def test_state_machine_creation_with_retry_in_constructor(sfn_client, sfn_role_arn, training_job_parameters):
591-
retry_state_name = "RetryStateName"
592-
all_fail_error = "Starts.ALL"
593-
interval_seconds = 1
594-
max_attempts = 2
595-
backoff_rate = 2
596-
task_resource = f"arn:{get_aws_partition()}:states:::sagemaker:createTrainingJob.sync"
597-
598-
# change the parameters to cause task state to fail
599-
training_job_parameters["AlgorithmSpecification"]["TrainingImage"] = "not_an_image"
600-
601535
asl_state_machine_definition = {
602536
"StartAt": retry_state_name,
603537
"States": {
@@ -608,7 +542,13 @@ def test_state_machine_creation_with_retry_in_constructor(sfn_client, sfn_role_a
608542
"End": True,
609543
"Retry": [
610544
{
611-
"ErrorEquals": [all_fail_error],
545+
"ErrorEquals": [timeout_error],
546+
"IntervalSeconds": interval_seconds,
547+
"MaxAttempts": max_attempts,
548+
"BackoffRate": backoff_rate
549+
},
550+
{
551+
"ErrorEquals": [task_failed_error],
612552
"IntervalSeconds": interval_seconds,
613553
"MaxAttempts": max_attempts,
614554
"BackoffRate": backoff_rate
@@ -618,22 +558,4 @@ def test_state_machine_creation_with_retry_in_constructor(sfn_client, sfn_role_a
618558
}
619559
}
620560

621-
task = steps.Task(
622-
retry_state_name,
623-
parameters=training_job_parameters,
624-
resource=task_resource,
625-
retry=steps.Retry(
626-
error_equals=[all_fail_error],
627-
interval_seconds=interval_seconds,
628-
max_attempts=max_attempts,
629-
backoff_rate=backoff_rate
630-
)
631-
)
632-
633-
workflow = Workflow(
634-
unique_name_from_base('Test_Retry_In_Constructor_Workflow'),
635-
definition=task,
636-
role=sfn_role_arn
637-
)
638-
639561
workflow_test_suite(sfn_client, workflow, asl_state_machine_definition, None)

0 commit comments

Comments
 (0)