@@ -83,3 +83,123 @@ def test_excessive_retry_attempts_error(self, t_mock, retry_policy):
8383 retry_policy .sleep (HTTPResponse (status = 503 ))
8484 # Internally urllib3 calls the increment function generating a new instance for every retry
8585 retry_policy = retry_policy .increment ()
86+
87+ @pytest .fixture ()
88+ def server_directed_retry_policy (self ) -> DatabricksRetryPolicy :
89+ return DatabricksRetryPolicy (
90+ delay_min = 1 ,
91+ delay_max = 30 ,
92+ stop_after_attempts_count = 3 ,
93+ stop_after_attempts_duration = 900 ,
94+ delay_default = 2 ,
95+ force_dangerous_codes = [],
96+ server_directed_only = True ,
97+ )
98+
99+ def test_server_directed_only__retries_with_retry_after (
100+ self , server_directed_retry_policy
101+ ):
102+ """429 + Retry-After header → should retry"""
103+ server_directed_retry_policy ._retry_start_time = time .time ()
104+ server_directed_retry_policy .command_type = CommandType .OTHER
105+ should_retry , msg = server_directed_retry_policy .should_retry (
106+ "POST" , 429 , has_retry_after = True
107+ )
108+ assert should_retry is True
109+
110+ def test_server_directed_only__no_retry_without_retry_after (
111+ self , server_directed_retry_policy
112+ ):
113+ """429 without Retry-After header → no retry"""
114+ server_directed_retry_policy ._retry_start_time = time .time ()
115+ server_directed_retry_policy .command_type = CommandType .OTHER
116+ should_retry , msg = server_directed_retry_policy .should_retry (
117+ "POST" , 429 , has_retry_after = False
118+ )
119+ assert should_retry is False
120+ assert "server_directed_only" in msg
121+
122+ def test_server_directed_only__no_retry_503_without_header (
123+ self , server_directed_retry_policy
124+ ):
125+ """503 without Retry-After header → no retry"""
126+ server_directed_retry_policy ._retry_start_time = time .time ()
127+ server_directed_retry_policy .command_type = CommandType .OTHER
128+ should_retry , msg = server_directed_retry_policy .should_retry (
129+ "POST" , 503 , has_retry_after = False
130+ )
131+ assert should_retry is False
132+ assert "server_directed_only" in msg
133+
134+ def test_server_directed_only__overrides_dangerous_codes (self ):
135+ """force_dangerous_codes=[500] + no Retry-After → no retry in server_directed_only mode"""
136+ policy = DatabricksRetryPolicy (
137+ delay_min = 1 ,
138+ delay_max = 30 ,
139+ stop_after_attempts_count = 3 ,
140+ stop_after_attempts_duration = 900 ,
141+ delay_default = 2 ,
142+ force_dangerous_codes = [500 ],
143+ server_directed_only = True ,
144+ )
145+ policy ._retry_start_time = time .time ()
146+ policy .command_type = CommandType .EXECUTE_STATEMENT
147+ should_retry , msg = policy .should_retry ("POST" , 500 , has_retry_after = False )
148+ assert should_retry is False
149+ assert "server_directed_only" in msg
150+
151+ def test_server_directed_only__non_retryable_codes_unaffected (
152+ self , server_directed_retry_policy
153+ ):
154+ """401/403/501 still don't retry even with Retry-After header"""
155+ server_directed_retry_policy ._retry_start_time = time .time ()
156+ server_directed_retry_policy .command_type = CommandType .OTHER
157+ for code in [401 , 403 , 501 ]:
158+ should_retry , msg = server_directed_retry_policy .should_retry (
159+ "POST" , code , has_retry_after = True
160+ )
161+ assert should_retry is False , f"Code { code }
162+
163+ def test_default_mode_unchanged (self , retry_policy ):
164+ """server_directed_only=False preserves existing behavior — 429 retries without header"""
165+ retry_policy ._retry_start_time = time .time ()
166+ retry_policy .command_type = CommandType .OTHER
167+ should_retry , msg = retry_policy .should_retry (
168+ "POST" , 429 , has_retry_after = False
169+ )
170+ assert should_retry is True
171+
172+ def test_server_directed_only__survives_new (self , server_directed_retry_policy ):
173+ """urllib3 calls .new() between retries to create a fresh policy instance.
174+ Verify that server_directed_only is carried over and still enforced."""
175+ server_directed_retry_policy ._retry_start_time = time .time ()
176+ server_directed_retry_policy .command_type = CommandType .OTHER
177+ new_policy = server_directed_retry_policy .new ()
178+ assert new_policy .server_directed_only is True
179+ # The new instance should still block retries without Retry-After
180+ should_retry , msg = new_policy .should_retry ("POST" , 429 , has_retry_after = False )
181+ assert should_retry is False
182+ assert "server_directed_only" in msg
183+
184+ def test_server_directed_only__execute_statement_with_retry_after (
185+ self , server_directed_retry_policy
186+ ):
187+ """EXECUTE_STATEMENT + 429 + Retry-After header → retry"""
188+ server_directed_retry_policy ._retry_start_time = time .time ()
189+ server_directed_retry_policy .command_type = CommandType .EXECUTE_STATEMENT
190+ should_retry , msg = server_directed_retry_policy .should_retry (
191+ "POST" , 429 , has_retry_after = True
192+ )
193+ assert should_retry is True
194+
195+ def test_404_does_not_retry_for_any_command_type (self , retry_policy ):
196+ """Test that 404 never retries for any CommandType"""
197+ retry_policy ._retry_start_time = time .time ()
198+
199+ # Test for each CommandType
200+ for command_type in CommandType :
201+ retry_policy .command_type = command_type
202+ should_retry , msg = retry_policy .should_retry ("POST" , 404 )
203+
204+ assert should_retry is False , f"404 should not retry for { command_type }
205+ assert "404" in msg or "NOT_FOUND" in msg
0 commit comments