Coverage for tests/test_server.py: 14%

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 

28from __future__ import annotations 

29 

30import asyncio 

31import os.path 

32import tempfile 

33import threading 

34import unittest 

35import unittest.mock 

36import uuid 

37from collections.abc import Generator, Iterator 

38from concurrent.futures import ThreadPoolExecutor 

39from contextlib import contextmanager 

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

41 

42from lsst.daf.butler import ( 

43 Butler, 

44 DataCoordinate, 

45 DatasetId, 

46 DatasetNotFoundError, 

47 DatasetRef, 

48 DatasetType, 

49 FileDataset, 

50 InvalidQueryError, 

51 LabeledButlerFactory, 

52 MissingDatasetTypeError, 

53 NoDefaultCollectionError, 

54 StorageClassFactory, 

55 UnknownComponentError, 

56) 

57from lsst.daf.butler.datastore import DatasetRefURIs 

58from lsst.daf.butler.registry import RegistryDefaults 

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

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

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

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

63from lsst.resources import ResourcePath 

64from lsst.resources.http import HttpResourcePath 

65 

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

67 import fastapi 

68 import httpx 

69 import safir.dependencies.logger 

70 from fastapi.testclient import TestClient 

71 

72 import lsst.daf.butler.remote_butler._get 

73 import lsst.daf.butler.remote_butler._query_results 

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

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

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

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

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

79 _EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY, 

80 RubinAuthenticationProvider, 

81 ) 

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

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

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

85 authorizer_dependency, 

86 butler_factory_dependency, 

87 ) 

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

89 from lsst.daf.butler.remote_butler.server.handlers._query_streaming import ( 

90 _stream_query_pages, 

91 ) 

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

93 from lsst.daf.butler.remote_butler.server_models import QueryCollectionsRequestModel, QueryKeepAliveModel 

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

95 

96 

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

98 

99 

100@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

101class ButlerClientServerTestCase(unittest.TestCase): 

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

103 

104 @classmethod 

105 def setUpClass(cls): 

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

107 cls.server_instance = server_instance 

108 cls.client = server_instance.client 

109 cls.butler = server_instance.remote_butler 

110 cls.butler_without_error_propagation = server_instance.remote_butler_without_error_propagation 

111 

112 cls.storageClassFactory = StorageClassFactory() 

113 

114 cls.repo = MetricTestRepo.create_from_butler( 

115 server_instance.direct_butler, server_instance.config_file_path 

116 ) 

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

118 cls.dataset_with_corrupted_data = _create_corrupted_dataset(cls.repo) 

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

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

121 cls.simple_dataset_ref = _create_simple_dataset(server_instance.direct_butler) 

122 

123 # Populate the test server. 

124 # The DatastoreMock is required because the datasets referenced in 

125 # these imports do not point at real files. 

126 direct_butler = server_instance.direct_butler 

127 DatastoreMock.apply(direct_butler) 

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

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

130 

131 def test_health_check(self): 

132 try: 

133 import importlib.metadata 

134 

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

136 except ModuleNotFoundError: 

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

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

139 self.assertEqual(response.status_code, 200) 

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

141 

142 def test_static_files(self): 

143 with tempfile.TemporaryDirectory() as tmpdir: 

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

145 fh.write("test data 123") 

146 

147 with mock_config() as server_config: 

148 server_config.static_files_path = tmpdir 

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

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

151 self.assertEqual(response.status_code, 200) 

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

153 

154 def test_dimension_universe(self): 

155 universe = self.butler.dimensions 

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

157 

158 def test_get_dataset_type(self): 

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

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

161 

162 with self.assertRaises(MissingDatasetTypeError): 

163 self.butler_without_error_propagation.get_dataset_type("not_bias") 

164 

165 def test_get_component_dataset_type(self): 

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

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

168 

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

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

171 parent definition. 

172 """ 

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

174 requested_paths: list[str] = [] 

175 original_get = self.butler._connection.get 

176 

177 def tracking_get(path, **kwargs): 

178 requested_paths.append(path) 

179 return original_get(path, **kwargs) 

180 

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

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

183 

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

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

186 for path in requested_paths: 

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

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

189 

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

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

192 

193 # An unknown component raises client-side. 

194 with self.assertRaises(KeyError): 

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

196 

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

198 with self.assertRaises(MissingDatasetTypeError): 

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

200 

201 def test_find_dataset(self): 

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

203 

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

205 self.assertIsInstance(ref, DatasetRef) 

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

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

208 

209 # Try again with variation of parameters. 

210 ref_new = self.butler.find_dataset( 

211 "bias", 

212 {"detector": 1}, 

213 collections="imported_g", 

214 instrument="Cam1", 

215 dimension_records=True, 

216 ) 

217 self.assertEqual(ref_new, ref) 

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

219 

220 ref_new = self.butler.find_dataset( 

221 ref.datasetType, 

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

223 collections="imported_g", 

224 storage_class=storage_class, 

225 ) 

226 self.assertEqual(ref_new, ref) 

227 

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

229 self.assertEqual(ref2, ref) 

230 

231 # Use detector name to find it. 

232 ref3 = self.butler.find_dataset( 

233 ref.datasetType, 

234 collections="imported_g", 

235 instrument="Cam1", 

236 full_name="Aa", 

237 ) 

238 self.assertEqual(ref2, ref3) 

239 

240 # Try expanded refs. 

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

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

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

244 

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

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

247 # code paths. 

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

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

250 

251 # Unknown dataset should not fail. 

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

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

254 

255 def test_instantiate_via_butler_http_search(self): 

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

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

258 RemoteButler pointing to the correct URL 

259 """ 

260 

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

262 # manipulations involving creating new ResourcePaths, and ResourcePath 

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

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

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

266 # 

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

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

269 # to the server under test. 

270 def override_read(http_resource_path): 

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

272 

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

274 

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

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

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

278 # server here. So mock it out. 

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

280 # Add access key to environment variables. RemoteButler 

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

282 # available. 

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

284 butler = Butler( 

285 server_url, 

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

287 run="collection2", 

288 ) 

289 self.enterContext(butler) 

290 self.assertIsInstance(butler, RemoteButler) 

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

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

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

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

295 self.assertFalse(butler._use_disabled_datastore_cache) 

296 

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

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

299 self.assertIsInstance(factory_created_butler, RemoteButler) 

300 self.assertTrue(factory_created_butler._use_disabled_datastore_cache) 

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

302 

303 def test_get(self): 

304 dataset_type = "test_metric_comp" 

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

306 collections = "ingest/run" 

307 # Test get() of a DatasetRef. 

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

309 metric = self.butler.get(ref) 

310 self.assertIsInstance(metric, MetricsExample) 

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

312 

313 # Test get() by DataId. 

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

315 self.assertEqual(metric, data_id_metric) 

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

317 kwarg_metric = self.butler.get( 

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

319 ) 

320 self.assertEqual(metric, kwarg_metric) 

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

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

323 kwarg_data_coordinate_metric = self.butler.get( 

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

325 ) 

326 self.assertEqual(metric, kwarg_data_coordinate_metric) 

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

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

329 with self.assertRaises(DatasetNotFoundError): 

330 self.butler_without_error_propagation.get( 

331 dataset_type, dataId=invalid_data_id, collections=collections 

332 ) 

333 

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

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

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

337 self.assertEqual(metric, default_collection_metric) 

338 

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

340 with self.assertRaises(NoDefaultCollectionError): 

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

342 

343 # Test looking up a non-existent ref 

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

345 with self.assertRaises(DatasetNotFoundError): 

346 self.butler_without_error_propagation.get(invalid_ref) 

347 

348 with self.assertRaises(RuntimeError): 

349 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data) 

350 

351 # Test storage class override 

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

353 

354 def check_sc_override(converted): 

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

356 self.assertIsInstance(converted, new_sc.pytype) 

357 self.assertEqual(metric, converted) 

358 

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

360 

361 # Test storage class override via DatasetRef. 

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

363 # Test storage class override via DatasetType. 

364 check_sc_override( 

365 self.butler.get( 

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

367 ) 

368 ) 

369 

370 # Test component override via DatasetRef. 

371 component_ref = ref.makeComponentRef("summary") 

372 component_data = self.butler.get(component_ref) 

373 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

374 

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

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

377 self.assertIsInstance(converted_component_data, DictConvertibleModel) 

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

379 

380 # Test component override via DatasetType. 

381 dataset_type_component_data = self.butler.get( 

382 component_ref.datasetType, component_ref.dataId, collections=collections 

383 ) 

384 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

385 

386 def test_get_formatter_receives_registry_ref(self): 

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

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

389 storage class override. 

390 

391 Formatters should read the write storage class from the 

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

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

394 """ 

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

396 

397 dataset_type = "test_metric_comp" 

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

399 collections = "ingest/run" 

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

401 write_storage_class = ref.datasetType.storageClass 

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

403 self.assertNotEqual(write_storage_class, override) 

404 

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

406 original = remote_get.generate_datastore_get_information 

407 

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

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

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

411 return original( 

412 fileLocations, registry_ref=registry_ref, read_ref=read_ref, parameters=parameters 

413 ) 

414 

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

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

417 

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

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

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

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

422 

423 def test_component_access_without_server_storage_class(self): 

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

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

426 

427 Storage classes for many dataset types are defined by science 

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

429 component handling must occur on the client. 

430 """ 

431 dataset_type = "test_metric_comp" 

432 component_type = "test_metric_comp.summary" 

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

434 collections = "ingest/run" 

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

436 component_ref = parent_ref.makeComponentRef("summary") 

437 

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

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

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

441 # parent's storage class. 

442 sent_dataset_types: list[str] = [] 

443 original_post = self.butler._connection.post 

444 

445 def tracking_post(path, model): 

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

447 if dataset_type_name is not None: 

448 sent_dataset_types.append(dataset_type_name) 

449 return original_post(path, model) 

450 

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

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

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

454 self.assertEqual(data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

455 

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

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

458 self.assertEqual(found, component_ref) 

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

460 

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

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

463 self.assertEqual(deferred_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

464 

465 # getDeferred() with a component DatasetRef. 

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

467 self.assertEqual(deferred_ref_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY) 

468 

469 # An empty component name raises rather than silently returning 

470 # the composite. 

471 with self.assertRaises(KeyError): 

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

473 with self.assertRaises(KeyError): 

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

475 

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

477 for name in sent_dataset_types: 

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

479 

480 def test_getURIs_no_components(self): 

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

482 def check_uri(uri: ResourcePath): 

483 self.assertIsNotNone(uris.primaryURI) 

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

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

486 

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

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

489 check_uri(uris.primaryURI) 

490 

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

492 

493 def test_getURIs_multiple_components(self): 

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

495 # URIs. 

496 dataset_type = "test_metric_comp" 

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

498 collections = "ingest/run" 

499 

500 def check_uris(uris: DatasetRefURIs): 

501 self.assertIsNone(uris.primaryURI) 

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

503 path = uris.componentURIs["summary"] 

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

505 data = path.read() 

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

507 

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

509 check_uris(uris) 

510 

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

512 with self.assertRaises(RuntimeError): 

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

514 

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

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

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

518 # behavior. 

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

520 componentRef = ref.makeComponentRef("summary") 

521 componentUris = self.butler.getURIs(componentRef) 

522 check_uris(componentUris) 

523 

524 def test_file_download_redirect(self): 

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

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

527 return self.client.get( 

528 uri, 

529 follow_redirects=False, 

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

531 ) 

532 

533 # Test behavior of a single-file dataset. 

534 response = get_download_redirect(self.simple_dataset_ref.id) 

535 self.assertEqual(response.status_code, 307) 

536 self.assertTrue(response.has_redirect_location) 

537 assert response.next_request is not None 

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

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

540 

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

542 self.assertEqual(response.status_code, 404) 

543 

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

545 dataset_type = "test_metric_comp" 

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

547 collections = "ingest/run" 

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

549 

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

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

552 self.assertEqual(response.status_code, 307) 

553 self.assertTrue(response.has_redirect_location) 

554 assert response.next_request is not None 

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

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

557 

558 # Unknown component. 

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

560 self.assertEqual(response.status_code, 404) 

561 

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

563 # composite". 

564 response = get_download_redirect(ref.id, None) 

565 self.assertEqual(response.status_code, 422) 

566 

567 # Unknown dataset. 

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

569 self.assertEqual(response.status_code, 404) 

570 

571 def test_auth_check(self): 

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

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

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

575 # handled by GafaelfawrIngress, not our app. 

576 with self.assertRaises(UnhandledServerError): 

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

578 

579 def test_exception_logging(self): 

580 app = create_app() 

581 

582 def raise_error(): 

583 raise RuntimeError("An unhandled error") 

584 

585 app.dependency_overrides[butler_factory_dependency] = raise_error 

586 client = TestClient(app, raise_server_exceptions=False) 

587 

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

589 mock_logger = NonCallableMock(["aerror"]) 

590 

591 async def noop(): 

592 pass 

593 

594 mock_logger.aerror.return_value = noop() 

595 

596 async def get_logger(): 

597 return mock_logger 

598 

599 mock_logger_dep.return_value = get_logger() 

600 client.get( 

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

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

603 ) 

604 mock_logger_dep.assert_called_once() 

605 

606 mock_logger.aerror.assert_called_once() 

607 args, kwargs = mock_logger.aerror.call_args 

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

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

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

611 

612 def test_query_keepalive(self): 

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

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

615 results. 

616 """ 

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

618 # immediately instead. 

619 with patch.object( 

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

621 ) as mock_timeout: 

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

623 # seen. 

624 with patch.object( 

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

626 ) as mock_keep_alive: 

627 mock_timeout.side_effect = _timeout_twice() 

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

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

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

631 self.assertGreaterEqual(mock_timeout.call_count, 3) 

632 self.assertGreaterEqual(mock_keep_alive.call_count, 2) 

633 

634 def test_query_retries(self): 

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

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

637 happens. 

638 """ 

639 query_event = threading.Event() 

640 retry_event = asyncio.Event() 

641 

642 async def block_first_request() -> None: 

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

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

645 # slot. 

646 query_event.set() 

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

648 # retry. 

649 await retry_event.wait() 

650 

651 async def block_second_request() -> None: 

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

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

654 retry_event.set() 

655 

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

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

658 

659 with ( 

660 patch.object( 

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

662 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES", 

663 new=1, 

664 ), 

665 patch.object( 

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

667 ), 

668 patch.object( 

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

670 "_block_query_for_unit_test", 

671 new=AsyncMock(wraps=block_first_request), 

672 ) as mock_first_client, 

673 patch.object( 

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

675 "_block_retry_for_unit_test", 

676 new=AsyncMock(wraps=block_second_request), 

677 ) as mock_second_client, 

678 ThreadPoolExecutor(max_workers=1) as exec1, 

679 ThreadPoolExecutor(max_workers=1) as exec2, 

680 ): 

681 first_butler = self.butler 

682 second_butler = self.butler.clone() 

683 

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

685 # query. 

686 future1 = exec1.submit(do_query, first_butler) 

687 event_reached = query_event.wait(60) 

688 if not event_reached: 

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

690 

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

692 # release the first client to finish its query. 

693 future2 = exec2.submit(do_query, second_butler) 

694 

695 result1 = future1.result(60) 

696 result2 = future2.result(60) 

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

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

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

700 # then the 2nd thread after it retries. 

701 self.assertEqual(mock_first_client.await_count, 2) 

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

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

704 # takes to finish. 

705 self.assertGreaterEqual(mock_second_client.await_count, 1) 

706 

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

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

709 def test_deprecated_collection_endpoints(self): 

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

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

712 json = self.butler._connection.get( 

713 "collection_info", 

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

715 ).json() 

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

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

718 

719 json = self.butler._connection.post( 

720 "query_collections", 

721 QueryCollectionsRequestModel( 

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

723 ), 

724 ).json() 

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

726 

727 def test_oversized_data_coordinate_upload(self): 

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

729 ref = self.simple_dataset_ref 

730 data_id = ref.dataId 

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

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

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

734 

735 

736@unittest.skipIf(not butler_server_is_available, butler_server_import_error) 

737class ButlerClientServerParityTestCase(unittest.TestCase): 

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

739 of component request and storage class override. 

740 

741 This uses its own server instance because 

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

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

744 """ 

745 

746 COMPONENT = "summary" 

747 

748 COMPOSITE_STORAGE_CLASS = "MetricsConversion" 

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

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

751 """ 

752 

753 COMPONENT_STORAGE_CLASS = "DictConvertibleModel" 

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

755 

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

757 COLLECTIONS = "ingest/run" 

758 

759 @classmethod 

760 def setUpClass(cls): 

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

762 cls.direct_butler = server_instance.direct_butler 

763 cls.remote_butler = server_instance.remote_butler 

764 cls.storageClassFactory = StorageClassFactory() 

765 repo = MetricTestRepo.create_from_butler( 

766 server_instance.direct_butler, server_instance.config_file_path