@@ -565,15 +565,19 @@ def test__upload_clienterror(project):
565565
566566def test__wait_for_registration_set_default (project ):
567567 mock_cfn_client = MagicMock (
568- spec = ["describe_type_registration" , "set_type_default_version" , "get_waiter" ]
568+ spec = ["describe_type_registration" , "set_type_default_version" ]
569569 )
570570 mock_cfn_client .describe_type_registration .return_value = (
571571 DESCRIBE_TYPE_COMPLETE_RETURN
572572 )
573573 mock_waiter = MagicMock (spec = ["wait" ])
574- mock_cfn_client .get_waiter .return_value = mock_waiter
575-
576- project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , True )
574+ patch_create_waiter = patch (
575+ "rpdk.core.project.create_waiter_with_client" ,
576+ autospec = True ,
577+ return_value = mock_waiter ,
578+ )
579+ with patch_create_waiter :
580+ project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , True )
577581
578582 mock_cfn_client .describe_type_registration .assert_called_once_with (
579583 RegistrationToken = REGISTRATION_TOKEN
@@ -586,7 +590,7 @@ def test__wait_for_registration_set_default(project):
586590
587591def test__wait_for_registration_set_default_fails (project ):
588592 mock_cfn_client = MagicMock (
589- spec = ["describe_type_registration" , "set_type_default_version" , "get_waiter" ]
593+ spec = ["describe_type_registration" , "set_type_default_version" ]
590594 )
591595 mock_cfn_client .describe_type_registration .return_value = (
592596 DESCRIBE_TYPE_COMPLETE_RETURN
@@ -595,9 +599,12 @@ def test__wait_for_registration_set_default_fails(project):
595599 BLANK_CLIENT_ERROR , "SetTypeDefaultVersion"
596600 )
597601 mock_waiter = MagicMock (spec = ["wait" ])
598- mock_cfn_client .get_waiter .return_value = mock_waiter
599-
600- with pytest .raises (DownstreamError ):
602+ patch_create_waiter = patch (
603+ "rpdk.core.project.create_waiter_with_client" ,
604+ autospec = True ,
605+ return_value = mock_waiter ,
606+ )
607+ with patch_create_waiter , pytest .raises (DownstreamError ):
601608 project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , True )
602609
603610 mock_cfn_client .describe_type_registration .assert_called_once_with (
@@ -611,15 +618,19 @@ def test__wait_for_registration_set_default_fails(project):
611618
612619def test__wait_for_registration_no_set_default (project ):
613620 mock_cfn_client = MagicMock (
614- spec = ["describe_type_registration" , "set_type_default_version" , "get_waiter" ]
621+ spec = ["describe_type_registration" , "set_type_default_version" ]
615622 )
616623 mock_cfn_client .describe_type_registration .return_value = (
617624 DESCRIBE_TYPE_COMPLETE_RETURN
618625 )
619626 mock_waiter = MagicMock (spec = ["wait" ])
620- mock_cfn_client .get_waiter .return_value = mock_waiter
621-
622- project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , False )
627+ patch_create_waiter = patch (
628+ "rpdk.core.project.create_waiter_with_client" ,
629+ autospec = True ,
630+ return_value = mock_waiter ,
631+ )
632+ with patch_create_waiter :
633+ project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , False )
623634
624635 mock_cfn_client .describe_type_registration .assert_called_once_with (
625636 RegistrationToken = REGISTRATION_TOKEN
@@ -630,7 +641,7 @@ def test__wait_for_registration_no_set_default(project):
630641
631642def test__wait_for_registration_waiter_fails (project ):
632643 mock_cfn_client = MagicMock (
633- spec = ["describe_type_registration" , "set_type_default_version" , "get_waiter" ]
644+ spec = ["describe_type_registration" , "set_type_default_version" ]
634645 )
635646 mock_cfn_client .describe_type_registration .return_value = (
636647 DESCRIBE_TYPE_FAILED_RETURN
@@ -641,9 +652,11 @@ def test__wait_for_registration_waiter_fails(project):
641652 "Waiter encountered a terminal failure state" ,
642653 DESCRIBE_TYPE_FAILED_RETURN ,
643654 )
644- mock_cfn_client .get_waiter .return_value = mock_waiter
645655
646- with pytest .raises (DownstreamError ):
656+ patch_create_waiter = patch (
657+ "rpdk.core.project.create_waiter_with_client" , return_value = mock_waiter
658+ )
659+ with patch_create_waiter , pytest .raises (DownstreamError ):
647660 project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , True )
648661
649662 mock_cfn_client .describe_type_registration .assert_called_once_with (
@@ -655,7 +668,7 @@ def test__wait_for_registration_waiter_fails(project):
655668
656669def test__wait_for_registration_waiter_fails_describe_fails (project ):
657670 mock_cfn_client = MagicMock (
658- spec = ["describe_type_registration" , "set_type_default_version" , "get_waiter" ]
671+ spec = ["describe_type_registration" , "set_type_default_version" ]
659672 )
660673 mock_cfn_client .describe_type_registration .side_effect = ClientError (
661674 BLANK_CLIENT_ERROR , "DescribeTypeRegistration"
@@ -667,9 +680,10 @@ def test__wait_for_registration_waiter_fails_describe_fails(project):
667680 DESCRIBE_TYPE_FAILED_RETURN ,
668681 )
669682
670- mock_cfn_client .get_waiter .return_value = mock_waiter
671-
672- with pytest .raises (DownstreamError ):
683+ patch_create_waiter = patch (
684+ "rpdk.core.project.create_waiter_with_client" , return_value = mock_waiter
685+ )
686+ with patch_create_waiter , pytest .raises (DownstreamError ):
673687 project ._wait_for_registration (mock_cfn_client , REGISTRATION_TOKEN , False )
674688
675689 mock_cfn_client .describe_type_registration .assert_called_once_with (
0 commit comments