Coverage for tests/test_server.py: 12%

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28import asyncio 

29import os.path 

30import tempfile 

31import threading 

32import unittest 

33import unittest.mock 

34import uuid 

35from concurrent.futures import ThreadPoolExecutor 

36from unittest.mock import DEFAULT, AsyncMock, NonCallableMock, patch 

37 

38from lsst.daf.butler import ( 

39 Butler, 

40 DataCoordinate, 

41 DatasetId, 

42 DatasetNotFoundError, 

43 DatasetRef, 

44 DatasetType, 

45 FileDataset, 

46 InvalidQueryError, 

47 LabeledButlerFactory, 

48 MissingDatasetTypeError, 

49 NoDefaultCollectionError, 

50 StorageClassFactory, 

51) 

52from lsst.daf.butler.datastore import DatasetRefURIs 

53from lsst.daf.butler.registry import RegistryDefaults 

54from lsst.daf.butler.tests import DatastoreMock, addDatasetType 

55from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel 

56from lsst.daf.butler.tests.server_available import butler_server_import_error, butler_server_is_available 

57from lsst.daf.butler.tests.utils import MetricsExample, MetricTestRepo, mock_env 

58from lsst.resources import ResourcePath 

59from lsst.resources.http import HttpResourcePath 

60 

61if butler_server_is_available: 61 ↛ 62line 61 didn't jump to line 62 because the condition on line 61 was never true

62 import fastapi 

63 import httpx 

64 import safir.dependencies.logger 

65 from fastapi.testclient import TestClient 

66 

67 import lsst.daf.butler.remote_butler._query_results 

68 import lsst.daf.butler.remote_butler.server.handlers._query_limits 

69 import lsst.daf.butler.remote_butler.server.handlers._query_streaming 

70 from lsst.daf.butler.remote_butler import ButlerServerError, RemoteButler 

71 from lsst.daf.butler.remote_butler.authentication.cadc import CadcAuthenticationProvider 

72 from lsst.daf.butler.remote_butler.authentication.rubin import ( 

73 _EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY, 

74 RubinAuthenticationProvider, 

75 ) 

76 from lsst.daf.butler.remote_butler.server import create_app 

77 from lsst.daf.butler.remote_butler.server._config import mock_config 

78 from lsst.daf.butler.remote_butler.server._dependencies import ( 

79 authorizer_dependency, 

80 butler_factory_dependency, 

81 ) 

82 from lsst.daf.butler.remote_butler.server._gafaelfawr import MockGafaelfawrGroupAuthorizer 

83 from lsst.daf.butler.remote_butler.server.handlers._utils import generate_file_download_uri 

84 from lsst.daf.butler.remote_butler.server_models import QueryCollectionsRequestModel 

85 from lsst.daf.butler.tests.server import TEST_REPOSITORY_NAME, UnhandledServerError, create_test_server 

86 

87 

88TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

89 

90 

91@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

92class ButlerClientServerTestCase(unittest.TestCase): 

93 """Test for Butler client/server.""" 

94 

95 @classmethod 

96 def setUpClass(cls): 

97 server_instance = cls.enterClassContext(create_test_server(TESTDIR)) 

98 cls.server_instance = server_instance 

99 cls.client = server_instance.client 

100 cls.butler = server_instance.remote_butler 

101 cls.butler_without_error_propagation = server_instance.remote_butler_without_error_propagation 

102 

103 cls.storageClassFactory = StorageClassFactory() 

104 

105 cls.repo = MetricTestRepo.create_from_butler( 

106 server_instance.direct_butler, server_instance.config_file_path 

107 ) 

108 # Add a file with corrupted data for testing error conditions 

109 cls.dataset_with_corrupted_data = _create_corrupted_dataset(cls.repo) 

110 # All of the datasets that come with MetricTestRepo are disassembled 

111 # composites. Add a simple dataset for testing the common case. 

112 cls.simple_dataset_ref = _create_simple_dataset(server_instance.direct_butler) 

113 

114 # Populate the test server. 

115 # The DatastoreMock is required because the datasets referenced in 

116 # these imports do not point at real files. 

117 direct_butler = server_instance.direct_butler 

118 DatastoreMock.apply(direct_butler) 

119 direct_butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

120 direct_butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/datasets.yaml") 

121 

122 def test_health_check(self): 

123 try: 

124 import importlib.metadata 

125 

126 importlib.metadata.metadata("lsst.daf.butler") 

127 except ModuleNotFoundError: 

128 raise self.skipTest("Standard python package metadata not available. Butler not pip installed.") 

129 response = self.client.get("/") 

130 self.assertEqual(response.status_code, 200) 

131 self.assertEqual(response.json()["name"], "butler") 

132 

133 def test_static_files(self): 

134 with tempfile.TemporaryDirectory() as tmpdir: 

135 with open(os.path.join(tmpdir, "temp.txt"), "w") as fh: 

136 fh.write("test data 123") 

137 

138 with mock_config() as server_config: 

139 server_config.static_files_path = tmpdir 

140 with create_test_server(TESTDIR, server_config=server_config) as server: 

141 response = server.client.get("/api/butler/configs/temp.txt") 

142 self.assertEqual(response.status_code, 200) 

143 self.assertEqual(response.text, "test data 123") 

144 

145 def test_dimension_universe(self): 

146 universe = self.butler.dimensions 

147 self.assertEqual(universe.namespace, "daf_butler") 

148 

149 def test_get_dataset_type(self): 

150 bias_type = self.butler.get_dataset_type("bias") 

151 self.assertEqual(bias_type.name, "bias") 

152 

153 with self.assertRaises(MissingDatasetTypeError): 

154 self.butler_without_error_propagation.get_dataset_type("not_bias") 

155 

156 def test_get_component_dataset_type(self): 

157 """Test that retrieving a component dataset type does not require the 

158 server to know the parent's storage class (DM-55497). 

159 

160 Component dataset type names must never be sent to the server, so 

161 the component dataset type is constructed on the client from the 

162 parent definition. 

163 """ 

164 # Track the dataset type names requested from the server. 

165 requested_paths: list[str] = [] 

166 original_get = self.butler._connection.get 

167 

168 def tracking_get(path, **kwargs): 

169 requested_paths.append(path) 

170 return original_get(path, **kwargs) 

171 

172 with patch.object(self.butler._connection, "get", side_effect=tracking_get): 

173 component_type = self.butler.get_dataset_type("bias.image") 

174 

175 parent_type = self.butler.get_dataset_type("bias") 

176 self.assertEqual(component_type, parent_type.makeComponentDatasetType("image")) 

177 for path in requested_paths: 

178 if path.startswith("dataset_type/"): 

179 self.assertNotIn(".", path, f"Component dataset type name was sent to the server: {path!r}") 

180 

181 # A second call should be served from the client-side cache. 

182 self.assertEqual(self.butler.get_dataset_type("bias.image"), component_type) 

183 

184 # An unknown component raises client-side. 

185 with self.assertRaises(KeyError): 

186 self.butler.get_dataset_type("bias.not_a_component") 

187 

188 # An unknown parent dataset type still raises the standard error. 

189 with self.assertRaises(MissingDatasetTypeError): 

190 self.butler_without_error_propagation.get_dataset_type("not_bias.image") 

191 

192 def test_find_dataset(self): 

193 storage_class = self.storageClassFactory.getStorageClass("Exposure") 

194 

195 ref = self.butler.find_dataset("bias", collections="imported_g", detector=1, instrument="Cam1") 

196 self.assertIsInstance(ref, DatasetRef) 

197 self.assertEqual(ref.id, uuid.UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22")) 

198 self.assertFalse(ref.dataId.hasRecords()) 

199 

200 # Try again with variation of parameters. 

201 ref_new = self.butler.find_dataset( 

202 "bias", 

203 {"detector": 1}, 

204 collections="imported_g", 

205 instrument="Cam1", 

206 dimension_records=True, 

207 ) 

208 self.assertEqual(ref_new, ref) 

209 self.assertTrue(ref_new.dataId.hasRecords()) 

210 

211 ref_new = self.butler.find_dataset( 

212 ref.datasetType, 

213 DataCoordinate.standardize(detector=1, instrument="Cam1", universe=self.butler.dimensions), 

214 collections="imported_g", 

215 storage_class=storage_class, 

216 ) 

217 self.assertEqual(ref_new, ref) 

218 

219 ref2 = self.butler.get_dataset(ref.id) 

220 self.assertEqual(ref2, ref) 

221 

222 # Use detector name to find it. 

223 ref3 = self.butler.find_dataset( 

224 ref.datasetType, 

225 collections="imported_g", 

226 instrument="Cam1", 

227 full_name="Aa", 

228 ) 

229 self.assertEqual(ref2, ref3) 

230 

231 # Try expanded refs. 

232 self.assertFalse(ref.dataId.hasRecords()) 

233 expanded = self.butler.get_dataset(ref.id, dimension_records=True) 

234 self.assertTrue(expanded.dataId.hasRecords()) 

235 

236 # The test datasets are all Exposure so storage class conversion 

237 # can not be tested until we fix that. For now at least test the 

238 # code paths. 

239 bias = self.butler.get_dataset(ref.id, storage_class=storage_class) 

240 self.assertEqual(bias.datasetType.storageClass, storage_class) 

241 

242 # Unknown dataset should not fail. 

243 self.assertIsNone(self.butler.get_dataset(uuid.uuid4())) 

244 self.assertIsNone(self.butler.get_dataset(uuid.uuid4(), storage_class="NumpyArray")) 

245 

246 def test_instantiate_via_butler_http_search(self): 

247 """Ensure that the primary Butler constructor's automatic search logic 

248 correctly locates and reads the configuration file and ends up with a 

249 RemoteButler pointing to the correct URL 

250 """ 

251 

252 # This is kind of a fragile test. Butler's search logic does a lot of 

253 # manipulations involving creating new ResourcePaths, and ResourcePath 

254 # doesn't use httpx so we can't easily inject the TestClient in there. 

255 # We don't have an actual valid HTTP URL to give to the constructor 

256 # because the test instance of the server is accessed via ASGI. 

257 # 

258 # Instead we just monkeypatch the HTTPResourcePath 'read' method and 

259 # hope that all ResourcePath HTTP reads during construction are going 

260 # to the server under test. 

261 def override_read(http_resource_path): 

262 return self.client.get(http_resource_path.geturl()).content 

263 

264 server_url = f"https://test.example/api/butler/repo/{TEST_REPOSITORY_NAME}/" 

265 

266 with patch.object(HttpResourcePath, "read", override_read): 

267 # RegistryDefaults.finish() needs to download the dimension 

268 # universe from the server, which will fail because there is no 

269 # server here. So mock it out. 

270 with patch.object(RegistryDefaults, "finish"): 

271 # Add access key to environment variables. RemoteButler 

272 # instantiation will throw an error if access key is not 

273 # available. 

274 with mock_env({_EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY: "fake-access-token"}): 

275 butler = Butler( 

276 server_url, 

277 collections=["collection1", "collection2"], 

278 run="collection2", 

279 ) 

280 self.enterContext(butler) 

281 self.assertIsInstance(butler, RemoteButler) 

282 self.assertEqual(butler._connection.server_url, server_url) 

283 self.assertEqual(butler.collections.defaults, ("collection1", "collection2")) 

284 self.assertEqual(butler.run, "collection2") 

285 # A butler created this way uses the default cache config. 

286 self.assertFalse(butler._use_disabled_datastore_cache) 

287 

288 butler_factory = LabeledButlerFactory({"server": server_url}) 

289 factory_created_butler = butler_factory.create_butler(label="server", access_token="token") 

290 self.assertIsInstance(factory_created_butler, RemoteButler) 

291 self.assertTrue(factory_created_butler._use_disabled_datastore_cache) 

292 self.assertEqual(factory_created_butler._connection.server_url, server_url) 

293 

294 def test_get(self): 

295 dataset_type = "test_metric_comp" 

296 data_id = {"instrument": "DummyCamComp", "visit": 423} 

297 collections = "ingest/run" 

298 # Test get() of a DatasetRef. 

299 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections) 

300 metric = self.butler.get(ref) 

301 self.assertIsInstance(metric, MetricsExample) 

302 self.assertEqual(metric.summary, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

303 

304 # Test get() by DataId. 

305 data_id_metric = self.butler.get(dataset_type, dataId=data_id, collections=collections) 

306 self.assertEqual(metric, data_id_metric) 

307 # Test get() by DataId dict augmented with kwargs. 

308 kwarg_metric = self.butler.get( 

309 dataset_type, dataId={"instrument": "DummyCamComp"}, collections=collections, visit=423 

310 ) 

311 self.assertEqual(metric, kwarg_metric) 

312 # Test get() by DataId DataCoordinate augmented with kwargs. 

313 coordinate = DataCoordinate.make_empty(self.butler.dimensions) 

314 kwarg_data_coordinate_metric = self.butler.get( 

315 dataset_type, dataId=coordinate, collections=collections, instrument="DummyCamComp", visit=423 

316 ) 

317 self.assertEqual(metric, kwarg_data_coordinate_metric) 

318 # Test get() of a non-existent DataId. 

319 invalid_data_id = {"instrument": "NotAValidlInstrument", "visit": 423} 

320 with self.assertRaises(DatasetNotFoundError): 

321 self.butler_without_error_propagation.get( 

322 dataset_type, dataId=invalid_data_id, collections=collections 

323 ) 

324 

325 # Test get() by DataId with default collections. 

326 butler_with_default_collection = self.butler.clone(collections="ingest/run") 

327 default_collection_metric = butler_with_default_collection.get(dataset_type, dataId=data_id) 

328 self.assertEqual(metric, default_collection_metric) 

329 

330 # Test get() by DataId with no collections specified. 

331 with self.assertRaises(NoDefaultCollectionError): 

332 self.butler_without_error_propagation.get(dataset_type, dataId=data_id) 

333 

334 # Test looking up a non-existent ref 

335 invalid_ref = ref.replace(id=uuid.uuid4()) 

336 with self.assertRaises(DatasetNotFoundError): 

337 self.butler_without_error_propagation.get(invalid_ref) 

338 

339 with self.assertRaises(RuntimeError): 

340 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data) 

341 

342 # Test storage class override 

343 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion") 

344 

345 def check_sc_override(converted): 

346 self.assertNotEqual(type(metric), type(converted)) 

347 self.assertIsInstance(converted, new_sc.pytype) 

348 self.assertEqual(metric, converted) 

349 

350 check_sc_override(self.butler.get(ref, storageClass=new_sc)) 

351 

352 # Test storage class override via DatasetRef. 

353 check_sc_override(self.butler.get(ref.overrideStorageClass("MetricsConversion"))) 

354 # Test storage class override via DatasetType. 

355 check_sc_override( 

356 self.butler.get( 

357 ref.datasetType.overrideStorageClass(new_sc), dataId=data_id, collections=collections 

358 ) 

359 ) 

360 

361 # Test component override via DatasetRef. 

362 component_ref = ref.makeComponentRef("summary") 

363 component_data = self.butler.get(component_ref) 

364 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

365 

366 # Test overriding both storage class and component via DatasetRef. 

367 converted_component_data = self.butler.get(component_ref, storageClass="DictConvertibleModel") 

368 self.assertIsInstance(converted_component_data, DictConvertibleModel) 

369 self.assertEqual(converted_component_data.content, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

370 

371 # Test component override via DatasetType. 

372 dataset_type_component_data = self.butler.get( 

373 component_ref.datasetType, component_ref.dataId, collections=collections 

374 ) 

375 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

376 

377 def test_component_access_without_server_storage_class(self): 

378 """Test that component dataset access via dataset type name does not 

379 require the server to know the parent's storage class. 

380 

381 Storage classes for many dataset types are defined by science 

382 pipelines packages that are only installed on the client, so all 

383 component handling must occur on the client. 

384 """ 

385 dataset_type = "test_metric_comp" 

386 component_type = "test_metric_comp.summary" 

387 data_id = {"instrument": "DummyCamComp", "visit": 423} 

388 collections = "ingest/run" 

389 parent_ref = self.butler.find_dataset(dataset_type, data_id, collections=collections) 

390 component_ref = parent_ref.makeComponentRef("summary") 

391 

392 # Track the dataset type names sent to the server. Component names 

393 # must never be sent, because converting a component dataset type 

394 # name to a DatasetType requires the server to instantiate the 

395 # parent's storage class. 

396 sent_dataset_types: list[str] = [] 

397 original_post = self.butler._connection.post 

398 

399 def tracking_post(path, model): 

400 dataset_type_name = getattr(model, "dataset_type", None) 

401 if dataset_type_name is not None: 

402 sent_dataset_types.append(dataset_type_name) 

403 return original_post(path, model) 

404 

405 with patch.object(self.butler._connection, "post", side_effect=tracking_post): 

406 # get() with a component dataset type name. 

407 data = self.butler.get(component_type, dataId=data_id, collections=collections) 

408 self.assertEqual(data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

409 

410 # find_dataset() with a component dataset type name. 

411 found = self.butler.find_dataset(component_type, data_id, collections=collections) 

412 self.assertEqual(found, component_ref) 

413 self.assertEqual(found.datasetType, component_ref.datasetType) 

414 

415 # getDeferred() with a component dataset type name. 

416 deferred_data = self.butler.getDeferred(component_type, data_id, collections=collections).get() 

417 self.assertEqual(deferred_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

418 

419 # getDeferred() with a component DatasetRef. 

420 deferred_ref_data = self.butler.getDeferred(component_ref).get() 

421 self.assertEqual(deferred_ref_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

422 

423 # An empty component name raises rather than silently returning 

424 # the composite. 

425 with self.assertRaises(KeyError): 

426 self.butler.get(f"{dataset_type}.", dataId=data_id, collections=collections) 

427 with self.assertRaises(KeyError): 

428 self.butler.getDeferred(f"{dataset_type}.", data_id, collections=collections) 

429 

430 self.assertGreater(len(sent_dataset_types), 0) 

431 for name in sent_dataset_types: 

432 self.assertNotIn(".", name, f"Component dataset type name {name!r} was sent to the server") 

433 

434 def test_getURIs_no_components(self): 

435 # This dataset does not have components, and should return one URI. 

436 def check_uri(uri: ResourcePath): 

437 self.assertIsNotNone(uris.primaryURI) 

438 self.assertEqual(uris.primaryURI.scheme, "https") 

439 self.assertEqual(uris.primaryURI.read(), b"123") 

440 

441 uris = self.butler.getURIs(self.simple_dataset_ref) 

442 self.assertEqual(len(uris.componentURIs), 0) 

443 check_uri(uris.primaryURI) 

444 

445 check_uri(self.butler.getURI(self.simple_dataset_ref)) 

446 

447 def test_getURIs_multiple_components(self): 

448 # This dataset has multiple components, so we should get back multiple 

449 # URIs. 

450 dataset_type = "test_metric_comp" 

451 data_id = {"instrument": "DummyCamComp", "visit": 423} 

452 collections = "ingest/run" 

453 

454 def check_uris(uris: DatasetRefURIs): 

455 self.assertIsNone(uris.primaryURI) 

456 self.assertEqual(len(uris.componentURIs), 3) 

457 path = uris.componentURIs["summary"] 

458 self.assertEqual(path.scheme, "https") 

459 data = path.read() 

460 self.assertEqual(data, b"AM1: 5.2\nAM2: 30.6\n") 

461 

462 uris = self.butler.getURIs(dataset_type, dataId=data_id, collections=collections) 

463 check_uris(uris) 

464 

465 # Calling getURI on a multi-file dataset raises an exception 

466 with self.assertRaises(RuntimeError): 

467 self.butler.getURI(dataset_type, dataId=data_id, collections=collections) 

468 

469 # getURIs does NOT respect component overrides on the DatasetRef, 

470 # instead returning the parent's URIs. Unclear if this is "correct" 

471 # from a conceptual point of view, but this matches DirectButler 

472 # behavior. 

473 ref = self.butler.find_dataset(dataset_type, data_id=data_id, collections=collections) 

474 componentRef = ref.makeComponentRef("summary") 

475 componentUris = self.butler.getURIs(componentRef) 

476 check_uris(componentUris) 

477 

478 def test_file_download_redirect(self): 

479 def get_download_redirect(id: DatasetId, component: str | None = None) -> httpx.Response: 

480 uri = generate_file_download_uri("http://unittest.test/", TEST_REPOSITORY_NAME, id, component) 

481 return self.client.get( 

482 uri, 

483 follow_redirects=False, 

484 headers=RubinAuthenticationProvider("mock-token").get_server_headers(), 

485 ) 

486 

487 # Test behavior of a single-file dataset. 

488 response = get_download_redirect(self.simple_dataset_ref.id) 

489 self.assertEqual(response.status_code, 307) 

490 self.assertTrue(response.has_redirect_location) 

491 assert response.next_request is not None 

492 self.assertEqual(response.next_request.url.scheme, "https") 

493 self.assertIn("test_int_DummyCamComp_R_d-r_423_ingest_run.json", response.next_request.url.path) 

494 

495 response = get_download_redirect(self.simple_dataset_ref.id, "somecomponent") 

496 self.assertEqual(response.status_code, 404) 

497 

498 # This dataset is a "disassembled composite" with multiple files. 

499 dataset_type = "test_metric_comp" 

500 data_id = {"instrument": "DummyCamComp", "visit": 423} 

501 collections = "ingest/run" 

502 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections) 

503 

504 # Getting single component of a multi-file "disassembled composite". 

505 response = get_download_redirect(ref.id, "summary") 

506 self.assertEqual(response.status_code, 307) 

507 self.assertTrue(response.has_redirect_location) 

508 assert response.next_request is not None 

509 self.assertEqual(response.next_request.url.scheme, "https") 

510 self.assertIn("test_metric_comp.summary", response.next_request.url.path) 

511 

512 # Unknown component. 

513 response = get_download_redirect(ref.id, "badcomponent") 

514 self.assertEqual(response.status_code, 404) 

515 

516 # Not specifying the component for a multi-file "disassembled 

517 # composite". 

518 response = get_download_redirect(ref.id, None) 

519 self.assertEqual(response.status_code, 422) 

520 

521 # Unknown dataset. 

522 response = get_download_redirect(uuid.UUID("59467c1b-fa13-4f7a-8ff8-cd83e092e563")) 

523 self.assertEqual(response.status_code, 404) 

524 

525 def test_auth_check(self): 

526 # This is checking that the unit-test middleware for validating the 

527 # authentication headers is working. It doesn't test actual server 

528 # functionality -- in a real deployment, the authentication headers are 

529 # handled by GafaelfawrIngress, not our app. 

530 with self.assertRaises(UnhandledServerError): 

531 self.client.get("/v1/dataset_type/int") 

532 

533 def test_exception_logging(self): 

534 app = create_app() 

535 

536 def raise_error(): 

537 raise RuntimeError("An unhandled error") 

538 

539 app.dependency_overrides[butler_factory_dependency] = raise_error 

540 client = TestClient(app, raise_server_exceptions=False) 

541 

542 with patch.object(safir.dependencies.logger, "logger_dependency") as mock_logger_dep: 

543 mock_logger = NonCallableMock(["aerror"]) 

544 

545 async def noop(): 

546 pass 

547 

548 mock_logger.aerror.return_value = noop() 

549 

550 async def get_logger(): 

551 return mock_logger 

552 

553 mock_logger_dep.return_value = get_logger() 

554 client.get( 

555 "/api/butler/repo/something/v1/dataset_type/int", 

556 headers={"X-Auth-Request-User": "user-name", "X-Butler-Client-Request-Id": "request-id"}, 

557 ) 

558 mock_logger_dep.assert_called_once() 

559 

560 mock_logger.aerror.assert_called_once() 

561 args, kwargs = mock_logger.aerror.call_args 

562 self.assertIsInstance(kwargs["exc_info"], RuntimeError) 

563 self.assertEqual(kwargs["clientRequestId"], "request-id") 

564 self.assertEqual(kwargs["user"], "user-name") 

565 

566 def test_query_keepalive(self): 

567 """Test that long-running queries stream keep-alive messages to stop 

568 the HTTP connection from closing before they are able to return 

569 results. 

570 """ 

571 # Normally it takes 15 seconds for a timeout -- mock it to trigger 

572 # immediately instead. 

573 with patch.object( 

574 lsst.daf.butler.remote_butler.server.handlers._query_streaming, "_timeout" 

575 ) as mock_timeout: 

576 # Hook into QueryDriver to track the number of keep-alives we have 

577 # seen. 

578 with patch.object( 

579 lsst.daf.butler.remote_butler._query_results, "_received_keep_alive" 

580 ) as mock_keep_alive: 

581 mock_timeout.side_effect = _timeout_twice() 

582 with self.butler.query() as query: 

583 datasets = list(query.datasets("bias", "imported_g")) 

584 self.assertEqual(len(datasets), 3) 

585 self.assertGreaterEqual(mock_timeout.call_count, 3) 

586 self.assertGreaterEqual(mock_keep_alive.call_count, 2) 

587 

588 def test_query_retries(self): 

589 """Test that the server will send HTTP status 503 to put backpressure 

590 on clients if it is overloaded, and that the client will retry if this 

591 happens. 

592 """ 

593 query_event = threading.Event() 

594 retry_event = asyncio.Event() 

595 

596 async def block_first_request() -> None: 

597 # Signal the unit tests that we have reached the critical section 

598 # in the server, where the first client has reserved the query 

599 # slot. 

600 query_event.set() 

601 # Block inside the query, until the 2nd client has been forced to 

602 # retry. 

603 await retry_event.wait() 

604 

605 async def block_second_request() -> None: 

606 # Release the first client, so it can finish its query and prevent 

607 # this client from being blocked on the next go-round. 

608 retry_event.set() 

609 

610 def do_query(butler: Butler) -> list[DatasetRef]: 

611 return butler.query_datasets("bias", "imported_g") 

612 

613 with ( 

614 patch.object( 

615 lsst.daf.butler.remote_butler.server.handlers._query_limits, 

616 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES", 

617 new=1, 

618 ), 

619 patch.object( 

620 lsst.daf.butler.remote_butler.server.handlers._query_limits, "_QUERY_RETRY_SECONDS", new=1 

621 ), 

622 patch.object( 

623 lsst.daf.butler.remote_butler.server.handlers._query_limits, 

624 "_block_query_for_unit_test", 

625 new=AsyncMock(wraps=block_first_request), 

626 ) as mock_first_client, 

627 patch.object( 

628 lsst.daf.butler.remote_butler.server.handlers._query_limits, 

629 "_block_retry_for_unit_test", 

630 new=AsyncMock(wraps=block_second_request), 

631 ) as mock_second_client, 

632 ThreadPoolExecutor(max_workers=1) as exec1, 

633 ThreadPoolExecutor(max_workers=1) as exec2, 

634 ): 

635 first_butler = self.butler 

636 second_butler = self.butler.clone() 

637 

638 # Run the first client up until the server starts executing its 

639 # query. 

640 future1 = exec1.submit(do_query, first_butler) 

641 event_reached = query_event.wait(60) 

642 if not event_reached: 

643 raise TimeoutError("Server did not execute query logic as expected.") 

644 

645 # Start the second client, which will trigger the retry logic and 

646 # release the first client to finish its query. 

647 future2 = exec2.submit(do_query, second_butler) 

648 

649 result1 = future1.result(60) 

650 result2 = future2.result(60) 

651 self.assertEqual(len(result1), 3) 

652 self.assertEqual(len(result2), 3) 

653 # The original thread should have gone through this section, and 

654 # then the 2nd thread after it retries. 

655 self.assertEqual(mock_first_client.await_count, 2) 

656 # We should have triggered the retry logic at least once, but it 

657 # might occur multiple times depending how long the first client 

658 # takes to finish. 

659 self.assertGreaterEqual(mock_second_client.await_count, 1) 

660 

661 # TODO DM-46204: This can be removed once the RSP recommended image has 

662 # been upgraded to a version that contains DM-46129. 

663 def test_deprecated_collection_endpoints(self): 

664 # These REST endpoints are no longer used by Butler client so they need 

665 # to be checked separately until they can be removed. 

666 json = self.butler._connection.get( 

667 "collection_info", 

668 params={"name": "imported_g", "include_doc": True, "include_parents": True}, 

669 ).json() 

670 self.assertEqual(json["name"], "imported_g") 

671 self.assertEqual(json["type"], 1) 

672 

673 json = self.butler._connection.post( 

674 "query_collections", 

675 QueryCollectionsRequestModel( 

676 search=["imported_*"], collection_types=[1], flatten_chains=False, include_chains=False 

677 ), 

678 ).json() 

679 self.assertCountEqual(json["collections"], ["imported_g", "imported_r"]) 

680 

681 def test_oversized_data_coordinate_upload(self): 

682 with self.butler.query() as query: 

683 ref = self.simple_dataset_ref 

684 data_id = ref.dataId 

685 data_coordinates = [DataCoordinate.standardize(data_id, visit=x) for x in range(100_001)] 

686 with self.assertRaisesRegex(InvalidQueryError, "data coordinate rows"): 

687 list(query.join_data_coordinates(data_coordinates).datasets(ref.datasetType, ref.run)) 

688 

689 

690@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

691class ButlerClientServerAuthorizationTestCase(unittest.TestCase): 

692 """Test authentication/authorization functionality.""" 

693 

694 def test_group_authorization(self): 

695 """Test that group membership repository authorization is checked when 

696 repository is accessed. 

697 """ 

698 with create_test_server(TESTDIR) as server_instance: 

699 mock = MockGafaelfawrGroupAuthorizer() 

700 server_instance.app.dependency_overrides[authorizer_dependency] = lambda: mock 

701 server_instance.direct_butler.registry.registerDatasetType( 

702 DatasetType("bias", [], "int", universe=server_instance.direct_butler.dimensions) 

703 ) 

704 server_instance.direct_butler.collections.register("collection") 

705 butler = server_instance.remote_butler 

706 mock.set_response(False) 

707 with self.assertRaises(ButlerServerError) as e: 

708 butler.get_dataset_type("bias") 

709 self.assertEqual(e.exception.status_code, 403) 

710 with self.assertRaises(ButlerServerError) as e: 

711 butler.query_datasets("bias", collections="*", find_first=False) 

712 self.assertEqual(e.exception.status_code, 403) 

713 

714 mock.set_response(True) 

715 self.assertEqual(butler.get_dataset_type("bias").name, "bias") 

716 self.assertEqual(butler.query_datasets("bias", collections="collection", explain=False), []) 

717 

718 def test_cadc_auth(self) -> None: 

719 """Test server running in CADC auth mode.""" 

720 with mock_config() as config: 

721 config.authentication = "cadc" 

722 config.gafaelfawr_url = "DISABLED" 

723 with create_test_server(TESTDIR, server_config=config) as instance: 

724 self.assertIsInstance(instance.remote_butler._connection.auth, CadcAuthenticationProvider) 

725 

726 # Set up a dataset backed by an HTTP URL. 

727 # CADC uses a plain HTTP service, not S3, for hosting Butler 

728 # artifacts. 

729 dataset_type = DatasetType("test", [], "int", universe=instance.direct_butler.dimensions) 

730 ref = DatasetRef( 

731 datasetType=dataset_type, 

732 dataId=DataCoordinate.makeEmpty(instance.direct_butler.dimensions), 

733 run="ingest/run", 

734 ) 

735 path = ResourcePath("https://fake-server.example/some-directory/file.json") 

736 dataset = FileDataset(path, ref) 

737 # ingest() insists on doing file existence checks, and we don't 

738 # have an HTTP server to point it at. 

739 with unittest.mock.patch( 

740 "lsst.daf.butler.datastores.fileDatastore.FileDatastore._standardizeIngestPath" 

741 ) as mock: 

742 mock.return_value = path 

743 instance.direct_butler.ingest(dataset, transfer="direct", record_validation_info=False) 

744 

745 # At the CADC, paths used for file download should NOT be a 

746 # signed URL, and should have authentication headers attached. 

747 def check_path(path_to_check: ResourcePath): 

748 self.assertEqual(str(path_to_check), str(path)) 

749 assert isinstance(path_to_check, HttpResourcePath) 

750 self.assertIsNotNone(path_to_check._extra_headers) 

751 self.assertIsNotNone(path_to_check._extra_headers.get("Authorization")) 

752 

753 check_path(instance.remote_butler.getURI(ref))