33
44package  com .azure .identity .implementation ;
55
6- import  com .azure .core .http .HttpResponse ;
76import  com .azure .core .test .http .MockHttpResponse ;
87import  org .junit .jupiter .api .Assertions ;
98import  org .junit .jupiter .api .Test ;
9+ import  org .junit .jupiter .params .ParameterizedTest ;
10+ import  org .junit .jupiter .params .provider .Arguments ;
11+ import  org .junit .jupiter .params .provider .MethodSource ;
1012
1113import  java .io .IOException ;
14+ import  java .util .stream .Stream ;
1215import  java .util .Arrays ;
1316import  java .util .LinkedList ;
1417import  java .util .Queue ;
@@ -34,29 +37,32 @@ public void testIMDSRetry() {
3437        }
3538    }
3639
37-     @ Test 
38-     public  void  testShouldRetry () {
40+     @ ParameterizedTest 
41+     @ MethodSource ("shouldRetryInDifferentScenarios" )
42+     public  void  testShouldRetry (String  headerValue , int  statusCode , boolean  expectedRetry , String  description ) {
3943        ImdsRetryStrategy  imdsRetryStrategy  = new  ImdsRetryStrategy ();
4044
41-         HttpResponse  httpResponse  = new  MockHttpResponse (null , 400 );
42- 
43-         Assertions .assertFalse (imdsRetryStrategy .shouldRetry (httpResponse ),
44-             "Imds Retry Strategy should not retry on 400 status response." );
45- 
46-         Assertions .assertTrue (imdsRetryStrategy .shouldRetry (new  MockHttpResponse (null , 410 )),
47-             "Imds Retry Strategy should retry on 410 status response." );
48- 
49-         Assertions .assertTrue (imdsRetryStrategy .shouldRetry (new  MockHttpResponse (null , 429 )),
50-             "Imds Retry Strategy should retry on 429 status response." );
51- 
52-         Assertions .assertTrue (imdsRetryStrategy .shouldRetry (new  MockHttpResponse (null , 500 )),
53-             "Imds Retry Strategy should retry on 429 status response." );
45+         MockHttpResponse  httpResponse  = new  MockHttpResponse (null , statusCode );
46+         if  (headerValue  != null ) {
47+             httpResponse .addHeader ("ResponseMessage" , headerValue );
48+         }
5449
55-         Assertions .assertTrue ( imdsRetryStrategy .shouldRetry (new   MockHttpResponse ( null ,  599 )), 
56-              "Imds Retry Strategy should retry on 429 status response." ); 
50+         Assertions .assertEquals ( expectedRetry ,  imdsRetryStrategy .shouldRetry (httpResponse ),  description ); 
51+     } 
5752
58-         Assertions .assertTrue (imdsRetryStrategy .shouldRetry (new  MockHttpResponse (null , 404 )),
59-             "Imds Retry Strategy should retry on 429 status response." );
53+     private  static  Stream <Arguments > shouldRetryInDifferentScenarios () {
54+         return  Stream .of (Arguments .of (null , 400 , false , "Imds Retry Strategy should not retry on 400 status response" ),
55+             Arguments .of (null , 410 , true , "Imds Retry Strategy should retry on 410 status response" ),
56+             Arguments .of (null , 429 , true , "Imds Retry Strategy should retry on 429 status response" ),
57+             Arguments .of (null , 500 , true , "Imds Retry Strategy should retry on 500 status reponse" ),
58+             Arguments .of (null , 599 , true , "Imds Retry Strategy should retry on 599 status response" ),
59+             Arguments .of (null , 404 , true , "Imds Retry Strategy should retry on 404 status response" ),
60+             Arguments .of ("A socket operation was attempted to an unreachable" , 403 , true ,
61+                 "Imds Retry Strategy should retry on 403 with unreachable message" ),
62+             Arguments .of ("Access denied" , 403 , false ,
63+                 "Imds Retry Strategy should not retry on 403 with Access Denied message" ),
64+             Arguments .of (null , 403 , false ,
65+                 "Imds Retry Strategy should not retry on 403 with no ResponseMessage header" ));
6066    }
6167
6268    @ Test 
0 commit comments