|
19 | 19 | from unittest.mock import patch |
20 | 20 |
|
21 | 21 | from gravitino.client.gravitino_client import GravitinoClient |
22 | | -from gravitino.client.gravitino_metalake import GravitinoMetalake |
23 | 22 | from gravitino.dto.authorization.user_dto import UserDTO |
24 | 23 | from gravitino.dto.requests.user_add_request import UserAddRequest |
25 | 24 | from gravitino.dto.responses.remove_response import RemoveResponse |
@@ -187,56 +186,83 @@ class TestGravitinoClientUserDelegates(unittest.TestCase): |
187 | 186 | """Verify that GravitinoClient correctly delegates User operations.""" |
188 | 187 |
|
189 | 188 | def _make_client(self): |
190 | | - metalake = mock_base.mock_load_metalake() |
191 | 189 | client = GravitinoClient.__new__(GravitinoClient) |
192 | | - client._metalake = metalake |
193 | | - return client, metalake |
| 190 | + return client |
194 | 191 |
|
195 | 192 | def test_client_add_user(self): |
196 | | - client, metalake = self._make_client() |
| 193 | + client = self._make_client() |
197 | 194 | user = _build_user_dto() |
198 | 195 | mock_resp = mock_base.mock_http_response(UserResponse(0, user).to_json()) |
199 | | - with patch( |
200 | | - "gravitino.utils.http_client.HTTPClient.post", return_value=mock_resp |
| 196 | + with ( |
| 197 | + patch.object( |
| 198 | + GravitinoClient, |
| 199 | + "get_metalake", |
| 200 | + return_value=mock_base.mock_load_metalake(), |
| 201 | + ), |
| 202 | + patch( |
| 203 | + "gravitino.utils.http_client.HTTPClient.post", return_value=mock_resp |
| 204 | + ), |
201 | 205 | ): |
202 | 206 | result = client.add_user("alice") |
203 | 207 | self.assertEqual("alice", result.name()) |
204 | 208 |
|
205 | 209 | def test_client_get_user(self): |
206 | | - client, metalake = self._make_client() |
| 210 | + client = self._make_client() |
207 | 211 | user = _build_user_dto(roles=["r1"]) |
208 | 212 | mock_resp = mock_base.mock_http_response(UserResponse(0, user).to_json()) |
209 | | - with patch( |
210 | | - "gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp |
| 213 | + with ( |
| 214 | + patch.object( |
| 215 | + GravitinoClient, |
| 216 | + "get_metalake", |
| 217 | + return_value=mock_base.mock_load_metalake(), |
| 218 | + ), |
| 219 | + patch("gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp), |
211 | 220 | ): |
212 | 221 | result = client.get_user("alice") |
213 | 222 | self.assertEqual(["r1"], result.roles()) |
214 | 223 |
|
215 | 224 | def test_client_remove_user(self): |
216 | | - client, metalake = self._make_client() |
| 225 | + client = self._make_client() |
217 | 226 | mock_resp = mock_base.mock_http_response(RemoveResponse(0, True).to_json()) |
218 | | - with patch( |
219 | | - "gravitino.utils.http_client.HTTPClient.delete", return_value=mock_resp |
| 227 | + with ( |
| 228 | + patch.object( |
| 229 | + GravitinoClient, |
| 230 | + "get_metalake", |
| 231 | + return_value=mock_base.mock_load_metalake(), |
| 232 | + ), |
| 233 | + patch( |
| 234 | + "gravitino.utils.http_client.HTTPClient.delete", return_value=mock_resp |
| 235 | + ), |
220 | 236 | ): |
221 | 237 | self.assertTrue(client.remove_user("alice")) |
222 | 238 |
|
223 | 239 | def test_client_list_users(self): |
224 | | - client, metalake = self._make_client() |
| 240 | + client = self._make_client() |
225 | 241 | users = [_build_user_dto("alice"), _build_user_dto("bob")] |
226 | 242 | mock_resp = mock_base.mock_http_response(UserListResponse(0, users).to_json()) |
227 | | - with patch( |
228 | | - "gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp |
| 243 | + with ( |
| 244 | + patch.object( |
| 245 | + GravitinoClient, |
| 246 | + "get_metalake", |
| 247 | + return_value=mock_base.mock_load_metalake(), |
| 248 | + ), |
| 249 | + patch("gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp), |
229 | 250 | ): |
230 | 251 | result = client.list_users() |
231 | 252 | self.assertEqual(["alice", "bob"], [u.name() for u in result]) |
232 | 253 |
|
233 | 254 | def test_client_list_user_names(self): |
234 | | - client, metalake = self._make_client() |
| 255 | + client = self._make_client() |
235 | 256 | mock_resp = mock_base.mock_http_response( |
236 | 257 | UserNamesListResponse(0, ["alice", "bob"]).to_json() |
237 | 258 | ) |
238 | | - with patch( |
239 | | - "gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp |
| 259 | + with ( |
| 260 | + patch.object( |
| 261 | + GravitinoClient, |
| 262 | + "get_metalake", |
| 263 | + return_value=mock_base.mock_load_metalake(), |
| 264 | + ), |
| 265 | + patch("gravitino.utils.http_client.HTTPClient.get", return_value=mock_resp), |
240 | 266 | ): |
241 | 267 | result = client.list_user_names() |
242 | 268 | self.assertEqual(["alice", "bob"], result) |
0 commit comments