Coverage for tests/test_server.py: 13%

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 UnknownComponentError, 

52) 

53from lsst.daf.butler.datastore import DatasetRefURIs 

54from lsst.daf.butler.registry import RegistryDefaults 

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

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

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

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

59from lsst.resources import ResourcePath 

60from lsst.resources.http import HttpResourcePath 

61 

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

63 import fastapi 

64 import httpx 

65 import safir.dependencies.logger 

66 from fastapi.testclient import TestClient 

67 

68 import lsst.daf.butler.remote_butler._get 

69 import lsst.daf.butler.remote_butler._query_results 

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

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

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

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

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

75 _EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY, 

76 RubinAuthenticationProvider, 

77 ) 

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

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

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

81 authorizer_dependency, 

82 butler_factory_dependency, 

83 ) 

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

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

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

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

88 

89 

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

91 

92 

93@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

94class ButlerClientServerTestCase(unittest.TestCase): 

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

96 

97 @classmethod 

98 def setUpClass(cls): 

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

100 cls.server_instance = server_instance 

101 cls.client = server_instance.client 

102 cls.butler = server_instance.remote_butler 

103 cls.butler_without_error_propagation = server_instance.remote_butler_without_error_propagation 

104 

105 cls.storageClassFactory = StorageClassFactory() 

106 

107 cls.repo = MetricTestRepo.create_from_butler( 

108 server_instance.direct_butler, server_instance.config_file_path 

109 ) 

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

111 cls.dataset_with_corrupted_data = _create_corrupted_dataset(cls.repo) 

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

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

114 cls.simple_dataset_ref = _create_simple_dataset(server_instance.direct_butler) 

115 

116 # Populate the test server. 

117 # The DatastoreMock is required because the datasets referenced in 

118 # these imports do not point at real files. 

119 direct_butler = server_instance.direct_butler 

120 DatastoreMock.apply(direct_butler) 

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

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

123 

124 def test_health_check(self): 

125 try: 

126 import importlib.metadata 

127 

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

129 except ModuleNotFoundError: 

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

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

132 self.assertEqual(response.status_code, 200) 

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

134 

135 def test_static_files(self): 

136 with tempfile.TemporaryDirectory() as tmpdir: 

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

138 fh.write("test data 123") 

139 

140 with mock_config() as server_config: 

141 server_config.static_files_path = tmpdir 

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

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

144 self.assertEqual(response.status_code, 200) 

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

146 

147 def test_dimension_universe(self): 

148 universe = self.butler.dimensions 

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

150 

151 def test_get_dataset_type(self): 

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

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

154 

155 with self.assertRaises(MissingDatasetTypeError): 

156 self.butler_without_error_propagation.get_dataset_type("not_bias") 

157 

158 def test_get_component_dataset_type(self): 

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

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

161 

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

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

164 parent definition. 

165 """ 

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

167 requested_paths: list[str] = [] 

168 original_get = self.butler._connection.get 

169 

170 def tracking_get(path, **kwargs): 

171 requested_paths.append(path) 

172 return original_get(path, **kwargs) 

173 

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

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

176 

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

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

179 for path in requested_paths: 

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

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

182 

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

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

185 

186 # An unknown component raises client-side. 

187 with self.assertRaises(KeyError): 

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

189 

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

191 with self.assertRaises(MissingDatasetTypeError): 

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

193 

194 def test_find_dataset(self): 

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

196 

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

198 self.assertIsInstance(ref, DatasetRef) 

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

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

201 

202 # Try again with variation of parameters. 

203 ref_new = self.butler.find_dataset( 

204 "bias", 

205 {"detector": 1}, 

206 collections="imported_g", 

207 instrument="Cam1", 

208 dimension_records=True, 

209 ) 

210 self.assertEqual(ref_new, ref) 

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

212 

213 ref_new = self.butler.find_dataset( 

214 ref.datasetType, 

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

216 collections="imported_g", 

217 storage_class=storage_class, 

218 ) 

219 self.assertEqual(ref_new, ref) 

220 

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

222 self.assertEqual(ref2, ref) 

223 

224 # Use detector name to find it. 

225 ref3 = self.butler.find_dataset( 

226 ref.datasetType, 

227 collections="imported_g", 

228 instrument="Cam1", 

229 full_name="Aa", 

230 ) 

231 self.assertEqual(ref2, ref3) 

232 

233 # Try expanded refs. 

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

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

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

237 

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

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

240 # code paths. 

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

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

243 

244 # Unknown dataset should not fail. 

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

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

247 

248 def test_instantiate_via_butler_http_search(self): 

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

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

251 RemoteButler pointing to the correct URL 

252 """ 

253 

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

255 # manipulations involving creating new ResourcePaths, and ResourcePath 

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

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

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

259 # 

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

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

262 # to the server under test. 

263 def override_read(http_resource_path): 

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

265 

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

267 

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

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

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

271 # server here. So mock it out. 

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

273 # Add access key to environment variables. RemoteButler 

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

275 # available. 

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

277 butler = Butler( 

278 server_url, 

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

280 run="collection2", 

281 ) 

282 self.enterContext(butler) 

283 self.assertIsInstance(butler, RemoteButler) 

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

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

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

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

288 self.assertFalse(butler._use_disabled_datastore_cache) 

289 

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

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

292 self.assertIsInstance(factory_created_butler, RemoteButler) 

293 self.assertTrue(factory_created_butler._use_disabled_datastore_cache) 

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

295 

296 def test_get(self): 

297 dataset_type = "test_metric_comp" 

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

299 collections = "ingest/run" 

300 # Test get() of a DatasetRef. 

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

302 metric = self.butler.get(ref) 

303 self.assertIsInstance(metric, MetricsExample) 

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

305 

306 # Test get() by DataId. 

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

308 self.assertEqual(metric, data_id_metric) 

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

310 kwarg_metric = self.butler.get( 

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

312 ) 

313 self.assertEqual(metric, kwarg_metric) 

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

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

316 kwarg_data_coordinate_metric = self.butler.get( 

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

318 ) 

319 self.assertEqual(metric, kwarg_data_coordinate_metric) 

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

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

322 with self.assertRaises(DatasetNotFoundError): 

323 self.butler_without_error_propagation.get( 

324 dataset_type, dataId=invalid_data_id, collections=collections 

325 ) 

326 

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

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

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

330 self.assertEqual(metric, default_collection_metric) 

331 

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

333 with self.assertRaises(NoDefaultCollectionError): 

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

335 

336 # Test looking up a non-existent ref 

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

338 with self.assertRaises(DatasetNotFoundError): 

339 self.butler_without_error_propagation.get(invalid_ref) 

340 

341 with self.assertRaises(RuntimeError): 

342 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data) 

343 

344 # Test storage class override 

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

346 

347 def check_sc_override(converted): 

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

349 self.assertIsInstance(converted, new_sc.pytype) 

350 self.assertEqual(metric, converted) 

351 

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

353 

354 # Test storage class override via DatasetRef. 

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

356 # Test storage class override via DatasetType. 

357 check_sc_override( 

358 self.butler.get( 

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

360 ) 

361 ) 

362 

363 # Test component override via DatasetRef. 

364 component_ref = ref.makeComponentRef("summary") 

365 component_data = self.butler.get(component_ref) 

366 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

367 

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

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

370 self.assertIsInstance(converted_component_data, DictConvertibleModel) 

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

372 

373 # Test component override via DatasetType. 

374 dataset_type_component_data = self.butler.get( 

375 component_ref.datasetType, component_ref.dataId, collections=collections 

376 ) 

377 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

378 

379 def test_get_formatter_receives_registry_ref(self): 

380 """The Formatter must always be given the `DatasetRef` as defined in 

381 the repository, so its storage class is unaffected by any read-time 

382 storage class override. 

383 

384 Formatters should read the write storage class from the 

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

386 hand the Formatter the same ref that `DirectButler` does. 

387 """ 

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

389 

390 dataset_type = "test_metric_comp" 

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

392 collections = "ingest/run" 

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

394 write_storage_class = ref.datasetType.storageClass 

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

396 self.assertNotEqual(write_storage_class, override) 

397 

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

399 original = remote_get.generate_datastore_get_information 

400 

401 def capturing(fileLocations, *, registry_ref, read_ref, parameters): 

402 captured["registry_storage_class"] = registry_ref.datasetType.storageClass 

403 captured["read_storage_class"] = read_ref.datasetType.storageClass 

404 return original( 

405 fileLocations, registry_ref=registry_ref, read_ref=read_ref, parameters=parameters 

406 ) 

407 

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

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

410 

411 # The ref given to the Formatter must carry the storage class from the 

412 # repository, and the override must instead travel via the read ref. 

413 self.assertEqual(captured["registry_storage_class"], write_storage_class) 

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

415 

416 def test_component_access_without_server_storage_class(self): 

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

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

419 

420 Storage classes for many dataset types are defined by science 

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

422 component handling must occur on the client. 

423 """ 

424 dataset_type = "test_metric_comp" 

425 component_type = "test_metric_comp.summary" 

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

427 collections = "ingest/run" 

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

429 component_ref = parent_ref.makeComponentRef("summary") 

430 

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

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

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

434 # parent's storage class. 

435 sent_dataset_types: list[str] = [] 

436 original_post = self.butler._connection.post 

437 

438 def tracking_post(path, model): 

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

440 if dataset_type_name is not None: 

441 sent_dataset_types.append(dataset_type_name) 

442 return original_post(path, model) 

443 

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

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

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

447 self.assertEqual(data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

448 

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

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

451 self.assertEqual(found, component_ref) 

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

453 

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

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

456 self.assertEqual(deferred_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

457 

458 # getDeferred() with a component DatasetRef. 

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

460 self.assertEqual(deferred_ref_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

461 

462 # An empty component name raises rather than silently returning 

463 # the composite. 

464 with self.assertRaises(KeyError): 

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

466 with self.assertRaises(KeyError): 

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

468 

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

470 for name in sent_dataset_types: 

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

472 

473 def test_getURIs_no_components(self): 

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

475 def check_uri(uri: ResourcePath): 

476 self.assertIsNotNone(uris.primaryURI) 

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

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

479 

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

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

482 check_uri(uris.primaryURI) 

483 

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

485 

486 def test_getURIs_multiple_components(self): 

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

488 # URIs. 

489 dataset_type = "test_metric_comp" 

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

491 collections = "ingest/run" 

492 

493 def check_uris(uris: DatasetRefURIs): 

494 self.assertIsNone(uris.primaryURI) 

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

496 path = uris.componentURIs["summary"] 

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

498 data = path.read() 

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

500 

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

502 check_uris(uris) 

503 

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

505 with self.assertRaises(RuntimeError): 

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

507 

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

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

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

511 # behavior. 

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

513 componentRef = ref.makeComponentRef("summary") 

514 componentUris = self.butler.getURIs(componentRef) 

515 check_uris(componentUris) 

516 

517 def test_file_download_redirect(self): 

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

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

520 return self.client.get( 

521 uri, 

522 follow_redirects=False, 

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

524 ) 

525 

526 # Test behavior of a single-file dataset. 

527 response = get_download_redirect(self.simple_dataset_ref.id) 

528 self.assertEqual(response.status_code, 307) 

529 self.assertTrue(response.has_redirect_location) 

530 assert response.next_request is not None 

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

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

533 

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

535 self.assertEqual(response.status_code, 404) 

536 

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

538 dataset_type = "test_metric_comp" 

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

540 collections = "ingest/run" 

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

542 

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

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

545 self.assertEqual(response.status_code, 307) 

546 self.assertTrue(response.has_redirect_location) 

547 assert response.next_request is not None 

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

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

550 

551 # Unknown component. 

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

553 self.assertEqual(response.status_code, 404) 

554 

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

556 # composite". 

557 response = get_download_redirect(ref.id, None) 

558 self.assertEqual(response.status_code, 422) 

559 

560 # Unknown dataset. 

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

562 self.assertEqual(response.status_code, 404) 

563 

564 def test_auth_check(self): 

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

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

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

568 # handled by GafaelfawrIngress, not our app. 

569 with self.assertRaises(UnhandledServerError): 

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

571 

572 def test_exception_logging(self): 

573 app = create_app() 

574 

575 def raise_error(): 

576 raise RuntimeError("An unhandled error") 

577 

578 app.dependency_overrides[butler_factory_dependency] = raise_error 

579 client = TestClient(app, raise_server_exceptions=False) 

580 

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

582 mock_logger = NonCallableMock(["aerror"]) 

583 

584 async def noop(): 

585 pass 

586 

587 mock_logger.aerror.return_value = noop() 

588 

589 async def get_logger(): 

590 return mock_logger 

591 

592 mock_logger_dep.return_value = get_logger() 

593 client.get( 

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

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

596 ) 

597 mock_logger_dep.assert_called_once() 

598 

599 mock_logger.aerror.assert_called_once() 

600 args, kwargs = mock_logger.aerror.call_args 

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

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

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

604 

605 def test_query_keepalive(self): 

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

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

608 results. 

609 """ 

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

611 # immediately instead. 

612 with patch.object( 

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

614 ) as mock_timeout: 

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

616 # seen. 

617 with patch.object( 

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

619 ) as mock_keep_alive: 

620 mock_timeout.side_effect = _timeout_twice() 

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

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

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

624 self.assertGreaterEqual(mock_timeout.call_count, 3) 

625 self.assertGreaterEqual(mock_keep_alive.call_count, 2) 

626 

627 def test_query_retries(self): 

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

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

630 happens. 

631 """ 

632 query_event = threading.Event() 

633 retry_event = asyncio.Event() 

634 

635 async def block_first_request() -> None: 

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

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

638 # slot. 

639 query_event.set() 

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

641 # retry. 

642 await retry_event.wait() 

643 

644 async def block_second_request() -> None: 

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

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

647 retry_event.set() 

648 

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

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

651 

652 with ( 

653 patch.object( 

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

655 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES", 

656 new=1, 

657 ), 

658 patch.object( 

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

660 ), 

661 patch.object( 

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

663 "_block_query_for_unit_test", 

664 new=AsyncMock(wraps=block_first_request), 

665 ) as mock_first_client, 

666 patch.object( 

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

668 "_block_retry_for_unit_test", 

669 new=AsyncMock(wraps=block_second_request), 

670 ) as mock_second_client, 

671 ThreadPoolExecutor(max_workers=1) as exec1, 

672 ThreadPoolExecutor(max_workers=1) as exec2, 

673 ): 

674 first_butler = self.butler 

675 second_butler = self.butler.clone() 

676 

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

678 # query. 

679 future1 = exec1.submit(do_query, first_butler) 

680 event_reached = query_event.wait(60) 

681 if not event_reached: 

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

683 

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

685 # release the first client to finish its query. 

686 future2 = exec2.submit(do_query, second_butler) 

687 

688 result1 = future1.result(60) 

689 result2 = future2.result(60) 

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

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

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

693 # then the 2nd thread after it retries. 

694 self.assertEqual(mock_first_client.await_count, 2) 

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

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

697 # takes to finish. 

698 self.assertGreaterEqual(mock_second_client.await_count, 1) 

699 

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

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

702 def test_deprecated_collection_endpoints(self): 

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

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

705 json = self.butler._connection.get( 

706 "collection_info", 

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

708 ).json() 

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

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

711 

712 json = self.butler._connection.post( 

713 "query_collections", 

714 QueryCollectionsRequestModel( 

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

716 ), 

717 ).json() 

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

719 

720 def test_oversized_data_coordinate_upload(self): 

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

722 ref = self.simple_dataset_ref 

723 data_id = ref.dataId 

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

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

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

727 

728 

729@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

730class ButlerClientServerParityTestCase(unittest.TestCase): 

731 """Test that `RemoteButler` and `DirectButler` agree for every combination 

732 of component request and storage class override. 

733 

734 This uses its own server instance because 

735 `ButlerClientServerTestCase` mocks the `DirectButler` datastore, which 

736 would leave nothing to compare the `RemoteButler` results against. 

737 """ 

738 

739 COMPONENT = "summary" 

740 

741 COMPOSITE_STORAGE_CLASS = "MetricsConversion" 

742 """Read storage class for the composite. Defines the same components as 

743 ``StructuredData`` but with a different Python type. 

744 """ 

745 

746 COMPONENT_STORAGE_CLASS = "DictConvertibleModel" 

747 """Read storage class for the component itself.""" 

748 

749 DATA_ID = {"instrument": "DummyCamComp", "visit": 423} 

750 COLLECTIONS = "ingest/run" 

751 

752 @classmethod 

753 def setUpClass(cls): 

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

755 cls.direct_butler = server_instance.direct_butler 

756 cls.remote_butler = server_instance.remote_butler 

757 cls.storageClassFactory = StorageClassFactory() 

758 repo = MetricTestRepo.create_from_butler( 

759 server_instance.direct_butler, server_instance.config_file_path 

760 ) 

761 # One dataset whose write storage class defines the component that will 

762 # be requested, and one whose write storage class defines no components 

763 # at all so that the component can only come from the read override. 

764 cls.refs = {} 

765 for name, storage_class in ( 

766 ("parity_with_components", "StructuredData"), 

767 ("parity_no_components", "StructuredDataNoComponents"), 

768 ):