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._get 

68 import lsst.daf.butler.remote_butler._query_results 

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

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

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

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

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

74 _EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY, 

75 RubinAuthenticationProvider, 

76 ) 

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

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

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

80 authorizer_dependency, 

81 butler_factory_dependency, 

82 ) 

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

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

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

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

87 

88 

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

90 

91 

92@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

93class ButlerClientServerTestCase(unittest.TestCase): 

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

95 

96 @classmethod 

97 def setUpClass(cls): 

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

99 cls.server_instance = server_instance 

100 cls.client = server_instance.client 

101 cls.butler = server_instance.remote_butler 

102 cls.butler_without_error_propagation = server_instance.remote_butler_without_error_propagation 

103 

104 cls.storageClassFactory = StorageClassFactory() 

105 

106 cls.repo = MetricTestRepo.create_from_butler( 

107 server_instance.direct_butler, server_instance.config_file_path 

108 ) 

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

110 cls.dataset_with_corrupted_data = _create_corrupted_dataset(cls.repo) 

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

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

113 cls.simple_dataset_ref = _create_simple_dataset(server_instance.direct_butler) 

114 

115 # Populate the test server. 

116 # The DatastoreMock is required because the datasets referenced in 

117 # these imports do not point at real files. 

118 direct_butler = server_instance.direct_butler 

119 DatastoreMock.apply(direct_butler) 

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

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

122 

123 def test_health_check(self): 

124 try: 

125 import importlib.metadata 

126 

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

128 except ModuleNotFoundError: 

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

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

131 self.assertEqual(response.status_code, 200) 

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

133 

134 def test_static_files(self): 

135 with tempfile.TemporaryDirectory() as tmpdir: 

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

137 fh.write("test data 123") 

138 

139 with mock_config() as server_config: 

140 server_config.static_files_path = tmpdir 

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

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

143 self.assertEqual(response.status_code, 200) 

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

145 

146 def test_dimension_universe(self): 

147 universe = self.butler.dimensions 

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

149 

150 def test_get_dataset_type(self): 

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

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

153 

154 with self.assertRaises(MissingDatasetTypeError): 

155 self.butler_without_error_propagation.get_dataset_type("not_bias") 

156 

157 def test_get_component_dataset_type(self): 

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

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

160 

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

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

163 parent definition. 

164 """ 

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

166 requested_paths: list[str] = [] 

167 original_get = self.butler._connection.get 

168 

169 def tracking_get(path, **kwargs): 

170 requested_paths.append(path) 

171 return original_get(path, **kwargs) 

172 

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

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

175 

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

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

178 for path in requested_paths: 

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

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

181 

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

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

184 

185 # An unknown component raises client-side. 

186 with self.assertRaises(KeyError): 

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

188 

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

190 with self.assertRaises(MissingDatasetTypeError): 

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

192 

193 def test_find_dataset(self): 

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

195 

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

197 self.assertIsInstance(ref, DatasetRef) 

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

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

200 

201 # Try again with variation of parameters. 

202 ref_new = self.butler.find_dataset( 

203 "bias", 

204 {"detector": 1}, 

205 collections="imported_g", 

206 instrument="Cam1", 

207 dimension_records=True, 

208 ) 

209 self.assertEqual(ref_new, ref) 

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

211 

212 ref_new = self.butler.find_dataset( 

213 ref.datasetType, 

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

215 collections="imported_g", 

216 storage_class=storage_class, 

217 ) 

218 self.assertEqual(ref_new, ref) 

219 

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

221 self.assertEqual(ref2, ref) 

222 

223 # Use detector name to find it. 

224 ref3 = self.butler.find_dataset( 

225 ref.datasetType, 

226 collections="imported_g", 

227 instrument="Cam1", 

228 full_name="Aa", 

229 ) 

230 self.assertEqual(ref2, ref3) 

231 

232 # Try expanded refs. 

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

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

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

236 

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

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

239 # code paths. 

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

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

242 

243 # Unknown dataset should not fail. 

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

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

246 

247 def test_instantiate_via_butler_http_search(self): 

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

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

250 RemoteButler pointing to the correct URL 

251 """ 

252 

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

254 # manipulations involving creating new ResourcePaths, and ResourcePath 

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

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

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

258 # 

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

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

261 # to the server under test. 

262 def override_read(http_resource_path): 

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

264 

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

266 

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

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

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

270 # server here. So mock it out. 

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

272 # Add access key to environment variables. RemoteButler 

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

274 # available. 

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

276 butler = Butler( 

277 server_url, 

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

279 run="collection2", 

280 ) 

281 self.enterContext(butler) 

282 self.assertIsInstance(butler, RemoteButler) 

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

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

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

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

287 self.assertFalse(butler._use_disabled_datastore_cache) 

288 

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

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

291 self.assertIsInstance(factory_created_butler, RemoteButler) 

292 self.assertTrue(factory_created_butler._use_disabled_datastore_cache) 

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

294 

295 def test_get(self): 

296 dataset_type = "test_metric_comp" 

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

298 collections = "ingest/run" 

299 # Test get() of a DatasetRef. 

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

301 metric = self.butler.get(ref) 

302 self.assertIsInstance(metric, MetricsExample) 

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

304 

305 # Test get() by DataId. 

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

307 self.assertEqual(metric, data_id_metric) 

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

309 kwarg_metric = self.butler.get( 

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

311 ) 

312 self.assertEqual(metric, kwarg_metric) 

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

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

315 kwarg_data_coordinate_metric = self.butler.get( 

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

317 ) 

318 self.assertEqual(metric, kwarg_data_coordinate_metric) 

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

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

321 with self.assertRaises(DatasetNotFoundError): 

322 self.butler_without_error_propagation.get( 

323 dataset_type, dataId=invalid_data_id, collections=collections 

324 ) 

325 

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

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

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

329 self.assertEqual(metric, default_collection_metric) 

330 

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

332 with self.assertRaises(NoDefaultCollectionError): 

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

334 

335 # Test looking up a non-existent ref 

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

337 with self.assertRaises(DatasetNotFoundError): 

338 self.butler_without_error_propagation.get(invalid_ref) 

339 

340 with self.assertRaises(RuntimeError): 

341 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data) 

342 

343 # Test storage class override 

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

345 

346 def check_sc_override(converted): 

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

348 self.assertIsInstance(converted, new_sc.pytype) 

349 self.assertEqual(metric, converted) 

350 

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

352 

353 # Test storage class override via DatasetRef. 

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

355 # Test storage class override via DatasetType. 

356 check_sc_override( 

357 self.butler.get( 

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

359 ) 

360 ) 

361 

362 # Test component override via DatasetRef. 

363 component_ref = ref.makeComponentRef("summary") 

364 component_data = self.butler.get(component_ref) 

365 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

366 

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

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

369 self.assertIsInstance(converted_component_data, DictConvertibleModel) 

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

371 

372 # Test component override via DatasetType. 

373 dataset_type_component_data = self.butler.get( 

374 component_ref.datasetType, component_ref.dataId, collections=collections 

375 ) 

376 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

377 

378 def test_get_formatter_receives_written_ref(self): 

379 """The Formatter must always be given the written `DatasetRef`, so its 

380 storage class matches the one the file was written with, regardless of 

381 any read-time storage class override. 

382 

383 Formatters should read the write storage class from the 

384 `FileDescriptor`, but some read it from the ref, so `RemoteButler` must 

385 hand the Formatter the same written ref that `DirectButler` does. 

386 """ 

387 remote_get = lsst.daf.butler.remote_butler._get 

388 

389 dataset_type = "test_metric_comp" 

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

391 collections = "ingest/run" 

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

393 write_storage_class = ref.datasetType.storageClass 

394 override = self.storageClassFactory.getStorageClass("MetricsConversion") 

395 self.assertNotEqual(write_storage_class, override) 

396 

397 captured: dict[str, object] = {} 

398 original = remote_get.generate_datastore_get_information 

399 

400 def capturing(fileLocations, *, ref, parameters, readStorageClass=None): 

401 captured["ref_storage_class"] = ref.datasetType.storageClass 

402 captured["read_storage_class"] = readStorageClass 

403 return original(fileLocations, ref=ref, parameters=parameters, readStorageClass=readStorageClass) 

404 

405 with patch.object(remote_get, "generate_datastore_get_information", side_effect=capturing): 

406 self.butler.get(ref, storageClass=override) 

407 

408 # The ref given to the Formatter must carry the write storage class, 

409 # and the override must instead travel via readStorageClass. 

410 self.assertEqual(captured["ref_storage_class"], write_storage_class) 

411 self.assertEqual(captured["read_storage_class"], override) 

412 

413 def test_component_access_without_server_storage_class(self): 

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

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

416 

417 Storage classes for many dataset types are defined by science 

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

419 component handling must occur on the client. 

420 """ 

421 dataset_type = "test_metric_comp" 

422 component_type = "test_metric_comp.summary" 

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

424 collections = "ingest/run" 

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

426 component_ref = parent_ref.makeComponentRef("summary") 

427 

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

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

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

431 # parent's storage class. 

432 sent_dataset_types: list[str] = [] 

433 original_post = self.butler._connection.post 

434 

435 def tracking_post(path, model): 

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

437 if dataset_type_name is not None: 

438 sent_dataset_types.append(dataset_type_name) 

439 return original_post(path, model) 

440 

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

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

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

444 self.assertEqual(data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

445 

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

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

448 self.assertEqual(found, component_ref) 

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

450 

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

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

453 self.assertEqual(deferred_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

454 

455 # getDeferred() with a component DatasetRef. 

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

457 self.assertEqual(deferred_ref_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

458 

459 # An empty component name raises rather than silently returning 

460 # the composite. 

461 with self.assertRaises(KeyError): 

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

463 with self.assertRaises(KeyError): 

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

465 

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

467 for name in sent_dataset_types: 

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

469 

470 def test_getURIs_no_components(self): 

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

472 def check_uri(uri: ResourcePath): 

473 self.assertIsNotNone(uris.primaryURI) 

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

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

476 

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

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

479 check_uri(uris.primaryURI) 

480 

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

482 

483 def test_getURIs_multiple_components(self): 

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

485 # URIs. 

486 dataset_type = "test_metric_comp" 

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

488 collections = "ingest/run" 

489 

490 def check_uris(uris: DatasetRefURIs): 

491 self.assertIsNone(uris.primaryURI) 

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

493 path = uris.componentURIs["summary"] 

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

495 data = path.read() 

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

497 

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

499 check_uris(uris) 

500 

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

502 with self.assertRaises(RuntimeError): 

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

504 

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

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

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

508 # behavior. 

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

510 componentRef = ref.makeComponentRef("summary") 

511 componentUris = self.butler.getURIs(componentRef) 

512 check_uris(componentUris) 

513 

514 def test_file_download_redirect(self): 

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

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

517 return self.client.get( 

518 uri, 

519 follow_redirects=False, 

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

521 ) 

522 

523 # Test behavior of a single-file dataset. 

524 response = get_download_redirect(self.simple_dataset_ref.id) 

525 self.assertEqual(response.status_code, 307) 

526 self.assertTrue(response.has_redirect_location) 

527 assert response.next_request is not None 

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

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

530 

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

532 self.assertEqual(response.status_code, 404) 

533 

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

535 dataset_type = "test_metric_comp" 

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

537 collections = "ingest/run" 

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

539 

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

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

542 self.assertEqual(response.status_code, 307) 

543 self.assertTrue(response.has_redirect_location) 

544 assert response.next_request is not None 

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

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

547 

548 # Unknown component. 

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

550 self.assertEqual(response.status_code, 404) 

551 

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

553 # composite". 

554 response = get_download_redirect(ref.id, None) 

555 self.assertEqual(response.status_code, 422) 

556 

557 # Unknown dataset. 

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

559 self.assertEqual(response.status_code, 404) 

560 

561 def test_auth_check(self): 

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

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

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

565 # handled by GafaelfawrIngress, not our app. 

566 with self.assertRaises(UnhandledServerError): 

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

568 

569 def test_exception_logging(self): 

570 app = create_app() 

571 

572 def raise_error(): 

573 raise RuntimeError("An unhandled error") 

574 

575 app.dependency_overrides[butler_factory_dependency] = raise_error 

576 client = TestClient(app, raise_server_exceptions=False) 

577 

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

579 mock_logger = NonCallableMock(["aerror"]) 

580 

581 async def noop(): 

582 pass 

583 

584 mock_logger.aerror.return_value = noop() 

585 

586 async def get_logger(): 

587 return mock_logger 

588 

589 mock_logger_dep.return_value = get_logger() 

590 client.get( 

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

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

593 ) 

594 mock_logger_dep.assert_called_once() 

595 

596 mock_logger.aerror.assert_called_once() 

597 args, kwargs = mock_logger.aerror.call_args 

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

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

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

601 

602 def test_query_keepalive(self): 

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

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

605 results. 

606 """ 

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

608 # immediately instead. 

609 with patch.object( 

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

611 ) as mock_timeout: 

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

613 # seen. 

614 with patch.object( 

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

616 ) as mock_keep_alive: 

617 mock_timeout.side_effect = _timeout_twice() 

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

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

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

621 self.assertGreaterEqual(mock_timeout.call_count, 3) 

622 self.assertGreaterEqual(mock_keep_alive.call_count, 2) 

623 

624 def test_query_retries(self): 

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

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

627 happens. 

628 """ 

629 query_event = threading.Event() 

630 retry_event = asyncio.Event() 

631 

632 async def block_first_request() -> None: 

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

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

635 # slot. 

636 query_event.set() 

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

638 # retry. 

639 await retry_event.wait() 

640 

641 async def block_second_request() -> None: 

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

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

644 retry_event.set() 

645 

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

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

648 

649 with ( 

650 patch.object( 

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

652 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES", 

653 new=1, 

654 ), 

655 patch.object( 

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

657 ), 

658 patch.object( 

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

660 "_block_query_for_unit_test", 

661 new=AsyncMock(wraps=block_first_request), 

662 ) as mock_first_client, 

663 patch.object( 

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

665 "_block_retry_for_unit_test", 

666 new=AsyncMock(wraps=block_second_request), 

667 ) as mock_second_client, 

668 ThreadPoolExecutor(max_workers=1) as exec1, 

669 ThreadPoolExecutor(max_workers=1) as exec2, 

670 ): 

671 first_butler = self.butler 

672 second_butler = self.butler.clone() 

673 

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

675 # query. 

676 future1 = exec1.submit(do_query, first_butler) 

677 event_reached = query_event.wait(60) 

678 if not event_reached: 

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

680 

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

682 # release the first client to finish its query. 

683 future2 = exec2.submit(do_query, second_butler) 

684 

685 result1 = future1.result(60) 

686 result2 = future2.result(60) 

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

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

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

690 # then the 2nd thread after it retries. 

691 self.assertEqual(mock_first_client.await_count, 2) 

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

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

694 # takes to finish. 

695 self.assertGreaterEqual(mock_second_client.await_count, 1) 

696 

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

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

699 def test_deprecated_collection_endpoints(self): 

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

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

702 json = self.butler._connection.get( 

703 "collection_info", 

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

705 ).json() 

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

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

708 

709 json = self.butler._connection.post( 

710 "query_collections", 

711 QueryCollectionsRequestModel( 

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

713 ), 

714 ).json() 

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

716 

717 def test_oversized_data_coordinate_upload(self): 

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

719 ref = self.simple_dataset_ref 

720 data_id = ref.dataId 

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

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

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

724 

725 

726@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

727class ButlerClientServerAuthorizationTestCase(unittest.TestCase): 

728 """Test authentication/authorization functionality.""" 

729 

730 def test_group_authorization(self): 

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

732 repository is accessed. 

733 """ 

734 with create_test_server(TESTDIR) as server_instance: 

735 mock = MockGafaelfawrGroupAuthorizer() 

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

737 server_instance.direct_butler.registry.registerDatasetType( 

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

739 ) 

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

741 butler = server_instance.remote_butler 

742 mock.set_response(False) 

743 with self.assertRaises(ButlerServerError) as e: 

744 butler.get_dataset_type("bias") 

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

746 with self.assertRaises(ButlerServerError) as e: 

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

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

749 

750 mock.set_response(True) 

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

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

753 

754 def test_cadc_auth(self) -> None: