@@ -121,8 +121,7 @@ def test_executor_without_query_tags_omits_from_api(self, mock_get_client):
121121 assert "query_tags" not in call_kwargs
122122
123123
124- def _make_warehouse (id , name , state , creator_name = "other@example.com" ,
125- enable_serverless_compute = False ):
124+ def _make_warehouse (id , name , state , creator_name = "other@example.com" , enable_serverless_compute = False ):
126125 """Helper to create a mock warehouse object."""
127126 w = mock .Mock ()
128127 w .id = id
@@ -141,33 +140,29 @@ class TestSortWithinTier:
141140 def test_serverless_first (self ):
142141 """Serverless warehouses should come before classic ones."""
143142 classic = _make_warehouse ("c1" , "Classic WH" , State .RUNNING )
144- serverless = _make_warehouse ("s1" , "Serverless WH" , State .RUNNING ,
145- enable_serverless_compute = True )
143+ serverless = _make_warehouse ("s1" , "Serverless WH" , State .RUNNING , enable_serverless_compute = True )
146144 result = _sort_within_tier ([classic , serverless ], current_user = None )
147145 assert result [0 ].id == "s1"
148146 assert result [1 ].id == "c1"
149147
150148 def test_serverless_before_user_owned (self ):
151149 """Serverless should be preferred over user-owned classic."""
152- classic_owned = _make_warehouse ("c1" , "My WH" , State .RUNNING ,
153- creator_name = "me@example.com" )
154- serverless_other = _make_warehouse ("s1" , "Other WH" , State .RUNNING ,
155- creator_name = "other@example.com" ,
156- enable_serverless_compute = True )
157- result = _sort_within_tier ([classic_owned , serverless_other ],
158- current_user = "me@example.com" )
150+ classic_owned = _make_warehouse ("c1" , "My WH" , State .RUNNING , creator_name = "me@example.com" )
151+ serverless_other = _make_warehouse (
152+ "s1" , "Other WH" , State .RUNNING , creator_name = "other@example.com" , enable_serverless_compute = True
153+ )
154+ result = _sort_within_tier ([classic_owned , serverless_other ], current_user = "me@example.com" )
159155 assert result [0 ].id == "s1"
160156
161157 def test_serverless_user_owned_first (self ):
162158 """Among serverless, user-owned should come first."""
163- serverless_other = _make_warehouse ("s1" , "Other Serverless" , State .RUNNING ,
164- creator_name = "other@example.com" ,
165- enable_serverless_compute = True )
166- serverless_owned = _make_warehouse ("s2" , "My Serverless" , State .RUNNING ,
167- creator_name = "me@example.com" ,
168- enable_serverless_compute = True )
169- result = _sort_within_tier ([serverless_other , serverless_owned ],
170- current_user = "me@example.com" )
159+ serverless_other = _make_warehouse (
160+ "s1" , "Other Serverless" , State .RUNNING , creator_name = "other@example.com" , enable_serverless_compute = True
161+ )
162+ serverless_owned = _make_warehouse (
163+ "s2" , "My Serverless" , State .RUNNING , creator_name = "me@example.com" , enable_serverless_compute = True
164+ )
165+ result = _sort_within_tier ([serverless_other , serverless_owned ], current_user = "me@example.com" )
171166 assert result [0 ].id == "s2"
172167 assert result [1 ].id == "s1"
173168
@@ -177,53 +172,46 @@ def test_empty_list(self):
177172 def test_no_current_user (self ):
178173 """Without a current user, only serverless preference applies."""
179174 classic = _make_warehouse ("c1" , "Classic" , State .RUNNING )
180- serverless = _make_warehouse ("s1" , "Serverless" , State .RUNNING ,
181- enable_serverless_compute = True )
175+ serverless = _make_warehouse ("s1" , "Serverless" , State .RUNNING , enable_serverless_compute = True )
182176 result = _sort_within_tier ([classic , serverless ], current_user = None )
183177 assert result [0 ].id == "s1"
184178
185179
186180class TestGetBestWarehouseServerless :
187181 """Tests for serverless preference in get_best_warehouse."""
188182
189- @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" ,
190- return_value = "me@example.com" )
183+ @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" , return_value = "me@example.com" )
191184 @mock .patch ("databricks_tools_core.sql.warehouse.get_workspace_client" )
192185 def test_prefers_serverless_within_running_shared (self , mock_client_fn , mock_user ):
193186 """Among running shared warehouses, serverless should be picked."""
194187 classic_shared = _make_warehouse ("c1" , "Shared WH" , State .RUNNING )
195- serverless_shared = _make_warehouse ("s1" , "Shared Serverless" , State .RUNNING ,
196- enable_serverless_compute = True )
188+ serverless_shared = _make_warehouse ("s1" , "Shared Serverless" , State .RUNNING , enable_serverless_compute = True )
197189 mock_client = mock .Mock ()
198190 mock_client .warehouses .list .return_value = [classic_shared , serverless_shared ]
199191 mock_client_fn .return_value = mock_client
200192
201193 result = get_best_warehouse ()
202194 assert result == "s1"
203195
204- @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" ,
205- return_value = "me@example.com" )
196+ @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" , return_value = "me@example.com" )
206197 @mock .patch ("databricks_tools_core.sql.warehouse.get_workspace_client" )
207198 def test_prefers_serverless_within_running_other (self , mock_client_fn , mock_user ):
208199 """Among running non-shared warehouses, serverless should be picked."""
209200 classic = _make_warehouse ("c1" , "My WH" , State .RUNNING )
210- serverless = _make_warehouse ("s1" , "Fast WH" , State .RUNNING ,
211- enable_serverless_compute = True )
201+ serverless = _make_warehouse ("s1" , "Fast WH" , State .RUNNING , enable_serverless_compute = True )
212202 mock_client = mock .Mock ()
213203 mock_client .warehouses .list .return_value = [classic , serverless ]
214204 mock_client_fn .return_value = mock_client
215205
216206 result = get_best_warehouse ()
217207 assert result == "s1"
218208
219- @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" ,
220- return_value = "me@example.com" )
209+ @mock .patch ("databricks_tools_core.sql.warehouse.get_current_username" , return_value = "me@example.com" )
221210 @mock .patch ("databricks_tools_core.sql.warehouse.get_workspace_client" )
222211 def test_tier_order_preserved_over_serverless (self , mock_client_fn , mock_user ):
223212 """A running shared classic should still beat a stopped serverless."""
224213 running_shared_classic = _make_warehouse ("c1" , "Shared WH" , State .RUNNING )
225- stopped_serverless = _make_warehouse ("s1" , "Fast WH" , State .STOPPED ,
226- enable_serverless_compute = True )
214+ stopped_serverless = _make_warehouse ("s1" , "Fast WH" , State .STOPPED , enable_serverless_compute = True )
227215 mock_client = mock .Mock ()
228216 mock_client .warehouses .list .return_value = [stopped_serverless , running_shared_classic ]
229217 mock_client_fn .return_value = mock_client
0 commit comments