Coverage for tests/test_server.py: 14%
529 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-09-09 02:00 -0700
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/>.
28from __future__ import annotations
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
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 DatasetTypeExpressionError, 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
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
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
97TESTDIR = os.path.abspath(os.path.dirname(__file__))
100@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
101class ButlerClientServerTestCase(unittest.TestCase):
102 """Test for Butler client/server."""
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
112 cls.storageClassFactory = StorageClassFactory()
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)
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")
131 def test_health_check(self):
132 try:
133 import importlib.metadata
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")
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")
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")
154 def test_dimension_universe(self):
155 universe = self.butler.dimensions
156 self.assertEqual(universe.namespace, "daf_butler")
158 def test_get_dataset_type(self):
159 bias_type = self.butler.get_dataset_type("bias")
160 self.assertEqual(bias_type.name, "bias")
162 with self.assertRaises(MissingDatasetTypeError):
163 self.butler_without_error_propagation.get_dataset_type("not_bias")
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).
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
177 def tracking_get(path, **kwargs):
178 requested_paths.append(path)
179 return original_get(path, **kwargs)
181 with patch.object(self.butler._connection, "get", side_effect=tracking_get):
182 component_type = self.butler.get_dataset_type("bias.image")
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}")
190 # A second call should be served from the client-side cache.
191 self.assertEqual(self.butler.get_dataset_type("bias.image"), component_type)
193 # An unknown component raises client-side.
194 with self.assertRaises(KeyError):
195 self.butler.get_dataset_type("bias.not_a_component")
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")
201 def test_get_dataset_type_with_invalid_name(self):
202 """Test that a syntactically invalid dataset type name is rejected on
203 the client without contacting the server (DM-53347).
204 """
205 with patch.object(self.butler._connection, "get") as mock:
206 with self.assertRaises(DatasetTypeExpressionError):
207 self.butler.get_dataset_type("...")
208 mock.assert_not_called()
210 def test_query_dataset_types_with_invalid_name(self):
211 """Test that a dataset type search expression the server cannot handle
212 is reported to the client as a user error rather than an internal
213 server error (DM-53347).
214 """
215 # "..." is a valid wildcard for the Ellipsis object but not as a
216 # string, where it parses as a component dataset type name.
217 with self.assertRaises(DatasetTypeExpressionError):
218 self.butler.registry.queryDatasetTypes("...")
220 with self.assertRaises(DatasetTypeExpressionError):
221 self.butler.registry.queryDatasetTypes("bias.image")
223 def test_find_dataset(self):
224 storage_class = self.storageClassFactory.getStorageClass("Exposure")
226 ref = self.butler.find_dataset("bias", collections="imported_g", detector=1, instrument="Cam1")
227 self.assertIsInstance(ref, DatasetRef)
228 self.assertEqual(ref.id, uuid.UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22"))
229 self.assertFalse(ref.dataId.hasRecords())
231 # Try again with variation of parameters.
232 ref_new = self.butler.find_dataset(
233 "bias",
234 {"detector": 1},
235 collections="imported_g",
236 instrument="Cam1",
237 dimension_records=True,
238 )
239 self.assertEqual(ref_new, ref)
240 self.assertTrue(ref_new.dataId.hasRecords())
242 ref_new = self.butler.find_dataset(
243 ref.datasetType,
244 DataCoordinate.standardize(detector=1, instrument="Cam1", universe=self.butler.dimensions),
245 collections="imported_g",
246 storage_class=storage_class,
247 )
248 self.assertEqual(ref_new, ref)
250 ref2 = self.butler.get_dataset(ref.id)
251 self.assertEqual(ref2, ref)
253 # Use detector name to find it.
254 ref3 = self.butler.find_dataset(
255 ref.datasetType,
256 collections="imported_g",
257 instrument="Cam1",
258 full_name="Aa",
259 )
260 self.assertEqual(ref2, ref3)
262 # Try expanded refs.
263 self.assertFalse(ref.dataId.hasRecords())
264 expanded = self.butler.get_dataset(ref.id, dimension_records=True)
265 self.assertTrue(expanded.dataId.hasRecords())
267 # The test datasets are all Exposure so storage class conversion
268 # can not be tested until we fix that. For now at least test the
269 # code paths.
270 bias = self.butler.get_dataset(ref.id, storage_class=storage_class)
271 self.assertEqual(bias.datasetType.storageClass, storage_class)
273 # Unknown dataset should not fail.
274 self.assertIsNone(self.butler.get_dataset(uuid.uuid4()))
275 self.assertIsNone(self.butler.get_dataset(uuid.uuid4(), storage_class="NumpyArray"))
277 def test_instantiate_via_butler_http_search(self):
278 """Ensure that the primary Butler constructor's automatic search logic
279 correctly locates and reads the configuration file and ends up with a
280 RemoteButler pointing to the correct URL
281 """
283 # This is kind of a fragile test. Butler's search logic does a lot of
284 # manipulations involving creating new ResourcePaths, and ResourcePath
285 # doesn't use httpx so we can't easily inject the TestClient in there.
286 # We don't have an actual valid HTTP URL to give to the constructor
287 # because the test instance of the server is accessed via ASGI.
288 #
289 # Instead we just monkeypatch the HTTPResourcePath 'read' method and
290 # hope that all ResourcePath HTTP reads during construction are going
291 # to the server under test.
292 def override_read(http_resource_path):
293 return self.client.get(http_resource_path.geturl()).content
295 server_url = f"https://test.example/api/butler/repo/{TEST_REPOSITORY_NAME}/"
297 with patch.object(HttpResourcePath, "read", override_read):
298 # RegistryDefaults.finish() needs to download the dimension
299 # universe from the server, which will fail because there is no
300 # server here. So mock it out.
301 with patch.object(RegistryDefaults, "finish"):
302 # Add access key to environment variables. RemoteButler
303 # instantiation will throw an error if access key is not
304 # available.
305 with mock_env({_EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY: "fake-access-token"}):
306 butler = Butler(
307 server_url,
308 collections=["collection1", "collection2"],
309 run="collection2",
310 )
311 self.enterContext(butler)
312 self.assertIsInstance(butler, RemoteButler)
313 self.assertEqual(butler._connection.server_url, server_url)
314 self.assertEqual(butler.collections.defaults, ("collection1", "collection2"))
315 self.assertEqual(butler.run, "collection2")
316 # A butler created this way uses the default cache config.
317 self.assertFalse(butler._use_disabled_datastore_cache)
319 butler_factory = LabeledButlerFactory({"server": server_url})
320 factory_created_butler = butler_factory.create_butler(label="server", access_token="token")
321 self.assertIsInstance(factory_created_butler, RemoteButler)
322 self.assertTrue(factory_created_butler._use_disabled_datastore_cache)
323 self.assertEqual(factory_created_butler._connection.server_url, server_url)
325 def test_get(self):
326 dataset_type = "test_metric_comp"
327 data_id = {"instrument": "DummyCamComp", "visit": 423}
328 collections = "ingest/run"
329 # Test get() of a DatasetRef.
330 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
331 metric = self.butler.get(ref)
332 self.assertIsInstance(metric, MetricsExample)
333 self.assertEqual(metric.summary, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
335 # Test get() by DataId.
336 data_id_metric = self.butler.get(dataset_type, dataId=data_id, collections=collections)
337 self.assertEqual(metric, data_id_metric)
338 # Test get() by DataId dict augmented with kwargs.
339 kwarg_metric = self.butler.get(
340 dataset_type, dataId={"instrument": "DummyCamComp"}, collections=collections, visit=423
341 )
342 self.assertEqual(metric, kwarg_metric)
343 # Test get() by DataId DataCoordinate augmented with kwargs.
344 coordinate = DataCoordinate.make_empty(self.butler.dimensions)
345 kwarg_data_coordinate_metric = self.butler.get(
346 dataset_type, dataId=coordinate, collections=collections, instrument="DummyCamComp", visit=423
347 )
348 self.assertEqual(metric, kwarg_data_coordinate_metric)
349 # Test get() of a non-existent DataId.
350 invalid_data_id = {"instrument": "NotAValidlInstrument", "visit": 423}
351 with self.assertRaises(DatasetNotFoundError):
352 self.butler_without_error_propagation.get(
353 dataset_type, dataId=invalid_data_id, collections=collections
354 )
356 # Test get() by DataId with default collections.
357 butler_with_default_collection = self.butler.clone(collections="ingest/run")
358 default_collection_metric = butler_with_default_collection.get(dataset_type, dataId=data_id)
359 self.assertEqual(metric, default_collection_metric)
361 # Test get() by DataId with no collections specified.
362 with self.assertRaises(NoDefaultCollectionError):
363 self.butler_without_error_propagation.get(dataset_type, dataId=data_id)
365 # Test looking up a non-existent ref
366 invalid_ref = ref.replace(id=uuid.uuid4())
367 with self.assertRaises(DatasetNotFoundError):
368 self.butler_without_error_propagation.get(invalid_ref)
370 with self.assertRaises(RuntimeError):
371 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data)
373 # Test storage class override
374 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
376 def check_sc_override(converted):
377 self.assertNotEqual(type(metric), type(converted))
378 self.assertIsInstance(converted, new_sc.pytype)
379 self.assertEqual(metric, converted)
381 check_sc_override(self.butler.get(ref, storageClass=new_sc))
383 # Test storage class override via DatasetRef.
384 check_sc_override(self.butler.get(ref.overrideStorageClass("MetricsConversion")))
385 # Test storage class override via DatasetType.
386 check_sc_override(
387 self.butler.get(
388 ref.datasetType.overrideStorageClass(new_sc), dataId=data_id, collections=collections
389 )
390 )
392 # Test component override via DatasetRef.
393 component_ref = ref.makeComponentRef("summary")
394 component_data = self.butler.get(component_ref)
395 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
397 # Test overriding both storage class and component via DatasetRef.
398 converted_component_data = self.butler.get(component_ref, storageClass="DictConvertibleModel")
399 self.assertIsInstance(converted_component_data, DictConvertibleModel)
400 self.assertEqual(converted_component_data.content, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
402 # Test component override via DatasetType.
403 dataset_type_component_data = self.butler.get(
404 component_ref.datasetType, component_ref.dataId, collections=collections
405 )
406 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
408 def test_get_formatter_receives_registry_ref(self):
409 """The Formatter must always be given the `DatasetRef` as defined in
410 the repository, so its storage class is unaffected by any read-time
411 storage class override.
413 Formatters should read the write storage class from the
414 `FileDescriptor`, but some read it from the ref, so `RemoteButler` must
415 hand the Formatter the same ref that `DirectButler` does.
416 """
417 remote_get = lsst.daf.butler.remote_butler._get
419 dataset_type = "test_metric_comp"
420 data_id = {"instrument": "DummyCamComp", "visit": 423}
421 collections = "ingest/run"
422 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
423 write_storage_class = ref.datasetType.storageClass
424 override = self.storageClassFactory.getStorageClass("MetricsConversion")
425 self.assertNotEqual(write_storage_class, override)
427 captured: dict[str, object] = {}
428 original = remote_get.generate_datastore_get_information
430 def capturing(fileLocations, *, registry_ref, read_ref, parameters):
431 captured["registry_storage_class"] = registry_ref.datasetType.storageClass
432 captured["read_storage_class"] = read_ref.datasetType.storageClass
433 return original(
434 fileLocations, registry_ref=registry_ref, read_ref=read_ref, parameters=parameters
435 )
437 with patch.object(remote_get, "generate_datastore_get_information", side_effect=capturing):
438 self.butler.get(ref, storageClass=override)
440 # The ref given to the Formatter must carry the storage class from the
441 # repository, and the override must instead travel via the read ref.
442 self.assertEqual(captured["registry_storage_class"], write_storage_class)
443 self.assertEqual(captured["read_storage_class"], override)
445 def test_component_access_without_server_storage_class(self):
446 """Test that component dataset access via dataset type name does not
447 require the server to know the parent's storage class.
449 Storage classes for many dataset types are defined by science
450 pipelines packages that are only installed on the client, so all
451 component handling must occur on the client.
452 """
453 dataset_type = "test_metric_comp"
454 component_type = "test_metric_comp.summary"
455 data_id = {"instrument": "DummyCamComp", "visit": 423}
456 collections = "ingest/run"
457 parent_ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
458 component_ref = parent_ref.makeComponentRef("summary")
460 # Track the dataset type names sent to the server. Component names
461 # must never be sent, because converting a component dataset type
462 # name to a DatasetType requires the server to instantiate the
463 # parent's storage class.
464 sent_dataset_types: list[str] = []
465 original_post = self.butler._connection.post
467 def tracking_post(path, model):
468 dataset_type_name = getattr(model, "dataset_type", None)
469 if dataset_type_name is not None:
470 sent_dataset_types.append(dataset_type_name)
471 return original_post(path, model)
473 with patch.object(self.butler._connection, "post", side_effect=tracking_post):
474 # get() with a component dataset type name.
475 data = self.butler.get(component_type, dataId=data_id, collections=collections)
476 self.assertEqual(data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
478 # find_dataset() with a component dataset type name.
479 found = self.butler.find_dataset(component_type, data_id, collections=collections)
480 self.assertEqual(found, component_ref)
481 self.assertEqual(found.datasetType, component_ref.datasetType)
483 # getDeferred() with a component dataset type name.
484 deferred_data = self.butler.getDeferred(component_type, data_id, collections=collections).get()
485 self.assertEqual(deferred_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
487 # getDeferred() with a component DatasetRef.
488 deferred_ref_data = self.butler.getDeferred(component_ref).get()
489 self.assertEqual(deferred_ref_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
491 # A trailing separator with no component name is not a valid
492 # dataset type name and is rejected rather than silently
493 # returning the composite.
494 with self.assertRaises(DatasetTypeExpressionError):
495 self.butler.get(f"{dataset_type}.", dataId=data_id, collections=collections)
496 with self.assertRaises(DatasetTypeExpressionError):
497 self.butler.getDeferred(f"{dataset_type}.", data_id, collections=collections)
499 self.assertGreater(len(sent_dataset_types), 0)
500 for name in sent_dataset_types:
501 self.assertNotIn(".", name, f"Component dataset type name {name!r} was sent to the server")
503 def test_getURIs_no_components(self):
504 # This dataset does not have components, and should return one URI.
505 def check_uri(uri: ResourcePath):
506 self.assertIsNotNone(uris.primaryURI)
507 self.assertEqual(uris.primaryURI.scheme, "https")
508 self.assertEqual(uris.primaryURI.read(), b"123")
510 uris = self.butler.getURIs(self.simple_dataset_ref)
511 self.assertEqual(len(uris.componentURIs), 0)
512 check_uri(uris.primaryURI)
514 check_uri(self.butler.getURI(self.simple_dataset_ref))
516 def test_getURIs_multiple_components(self):
517 # This dataset has multiple components, so we should get back multiple
518 # URIs.
519 dataset_type = "test_metric_comp"
520 data_id = {"instrument": "DummyCamComp", "visit": 423}
521 collections = "ingest/run"
523 def check_uris(uris: DatasetRefURIs):
524 self.assertIsNone(uris.primaryURI)
525 self.assertEqual(len(uris.componentURIs), 3)
526 path = uris.componentURIs["summary"]
527 self.assertEqual(path.scheme, "https")
528 data = path.read()
529 self.assertEqual(data, b"AM1: 5.2\nAM2: 30.6\n")
531 uris = self.butler.getURIs(dataset_type, dataId=data_id, collections=collections)
532 check_uris(uris)
534 # Calling getURI on a multi-file dataset raises an exception
535 with self.assertRaises(RuntimeError):
536 self.butler.getURI(dataset_type, dataId=data_id, collections=collections)
538 # getURIs does NOT respect component overrides on the DatasetRef,
539 # instead returning the parent's URIs. Unclear if this is "correct"
540 # from a conceptual point of view, but this matches DirectButler
541 # behavior.
542 ref = self.butler.find_dataset(dataset_type, data_id=data_id, collections=collections)
543 componentRef = ref.makeComponentRef("summary")
544 componentUris = self.butler.getURIs(componentRef)
545 check_uris(componentUris)
547 def test_file_download_redirect(self):
548 def get_download_redirect(id: DatasetId, component: str | None = None) -> httpx.Response:
549 uri = generate_file_download_uri("http://unittest.test/", TEST_REPOSITORY_NAME, id, component)
550 return self.client.get(
551 uri,
552 follow_redirects=False,
553 headers=RubinAuthenticationProvider("mock-token").get_server_headers(),
554 )
556 # Test behavior of a single-file dataset.
557 response = get_download_redirect(self.simple_dataset_ref.id)
558 self.assertEqual(response.status_code, 307)
559 self.assertTrue(response.has_redirect_location)
560 assert response.next_request is not None
561 self.assertEqual(response.next_request.url.scheme, "https")
562 self.assertIn("test_int_DummyCamComp_R_d-r_423_ingest_run.json", response.next_request.url.path)
564 response = get_download_redirect(self.simple_dataset_ref.id, "somecomponent")
565 self.assertEqual(response.status_code, 404)
567 # This dataset is a "disassembled composite" with multiple files.
568 dataset_type = "test_metric_comp"
569 data_id = {"instrument": "DummyCamComp", "visit": 423}
570 collections = "ingest/run"
571 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
573 # Getting single component of a multi-file "disassembled composite".
574 response = get_download_redirect(ref.id, "summary")
575 self.assertEqual(response.status_code, 307)
576 self.assertTrue(response.has_redirect_location)
577 assert response.next_request is not None
578 self.assertEqual(response.next_request.url.scheme, "https")
579 self.assertIn("test_metric_comp.summary", response.next_request.url.path)
581 # Unknown component.
582 response = get_download_redirect(ref.id, "badcomponent")
583 self.assertEqual(response.status_code, 404)
585 # Not specifying the component for a multi-file "disassembled
586 # composite".
587 response = get_download_redirect(ref.id, None)
588 self.assertEqual(response.status_code, 422)
590 # Unknown dataset.
591 response = get_download_redirect(uuid.UUID("59467c1b-fa13-4f7a-8ff8-cd83e092e563"))
592 self.assertEqual(response.status_code, 404)
594 def test_auth_check(self):
595 # This is checking that the unit-test middleware for validating the
596 # authentication headers is working. It doesn't test actual server
597 # functionality -- in a real deployment, the authentication headers are
598 # handled by GafaelfawrIngress, not our app.
599 with self.assertRaises(UnhandledServerError):
600 self.client.get("/v1/dataset_type/int")
602 def test_exception_logging(self):
603 app = create_app()
605 def raise_error():
606 raise RuntimeError("An unhandled error")
608 app.dependency_overrides[butler_factory_dependency] = raise_error
609 client = TestClient(app, raise_server_exceptions=False)
611 with patch.object(safir.dependencies.logger, "logger_dependency") as mock_logger_dep:
612 mock_logger = NonCallableMock(["aerror"])
614 async def noop():
615 pass
617 mock_logger.aerror.return_value = noop()
619 async def get_logger():
620 return mock_logger
622 mock_logger_dep.return_value = get_logger()
623 client.get(
624 "/api/butler/repo/something/v1/dataset_type/int",
625 headers={"X-Auth-Request-User": "user-name", "X-Butler-Client-Request-Id": "request-id"},
626 )
627 mock_logger_dep.assert_called_once()
629 mock_logger.aerror.assert_called_once()
630 args, kwargs = mock_logger.aerror.call_args
631 self.assertIsInstance(kwargs["exc_info"], RuntimeError)
632 self.assertEqual(kwargs["clientRequestId"], "request-id")
633 self.assertEqual(kwargs["user"], "user-name")
635 def test_query_keepalive(self):
636 """Test that long-running queries stream keep-alive messages to stop
637 the HTTP connection from closing before they are able to return
638 results.
639 """
640 # Normally it takes 15 seconds for a timeout -- mock it to trigger
641 # immediately instead.
642 with patch.object(
643 lsst.daf.butler.remote_butler.server.handlers._query_streaming, "_timeout"
644 ) as mock_timeout:
645 # Hook into QueryDriver to track the number of keep-alives we have
646 # seen.
647 with patch.object(
648 lsst.daf.butler.remote_butler._query_results, "_received_keep_alive"
649 ) as mock_keep_alive:
650 mock_timeout.side_effect = _timeout_twice()
651 with self.butler.query() as query:
652 datasets = list(query.datasets("bias", "imported_g"))
653 self.assertEqual(len(datasets), 3)
654 self.assertGreaterEqual(mock_timeout.call_count, 3)
655 self.assertGreaterEqual(mock_keep_alive.call_count, 2)
657 def test_query_retries(self):
658 """Test that the server will send HTTP status 503 to put backpressure
659 on clients if it is overloaded, and that the client will retry if this
660 happens.
661 """
662 query_event = threading.Event()
663 retry_event = asyncio.Event()
665 async def block_first_request() -> None:
666 # Signal the unit tests that we have reached the critical section
667 # in the server, where the first client has reserved the query
668 # slot.
669 query_event.set()
670 # Block inside the query, until the 2nd client has been forced to
671 # retry.
672 await retry_event.wait()
674 async def block_second_request() -> None:
675 # Release the first client, so it can finish its query and prevent
676 # this client from being blocked on the next go-round.
677 retry_event.set()
679 def do_query(butler: Butler) -> list[DatasetRef]:
680 return butler.query_datasets("bias", "imported_g")
682 with (
683 patch.object(
684 lsst.daf.butler.remote_butler.server.handlers._query_limits,
685 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES",
686 new=1,
687 ),
688 patch.object(
689 lsst.daf.butler.remote_butler.server.handlers._query_limits, "_QUERY_RETRY_SECONDS", new=1
690 ),
691 patch.object(
692 lsst.daf.butler.remote_butler.server.handlers._query_limits,
693 "_block_query_for_unit_test",
694 new=AsyncMock(wraps=block_first_request),
695 ) as mock_first_client,
696 patch.object(
697 lsst.daf.butler.remote_butler.server.handlers._query_limits,
698 "_block_retry_for_unit_test",
699 new=AsyncMock(wraps=block_second_request),
700 ) as mock_second_client,
701 ThreadPoolExecutor(max_workers=1) as exec1,
702 ThreadPoolExecutor(max_workers=1) as exec2,
703 ):
704 first_butler = self.butler
705 second_butler = self.butler.clone()
707 # Run the first client up until the server starts executing its
708 # query.
709 future1 = exec1.submit(do_query, first_butler)
710 event_reached = query_event.wait(60)
711 if not event_reached:
712 raise TimeoutError("Server did not execute query logic as expected.")
714 # Start the second client, which will trigger the retry logic and
715 # release the first client to finish its query.
716 future2 = exec2.submit(do_query, second_butler)
718 result1 = future1.result(60)
719 result2 = future2.result(60)
720 self.assertEqual(len(result1), 3)
721 self.assertEqual(len(result2), 3)
722 # The original thread should have gone through this section, and
723 # then the 2nd thread after it retries.
724 self.assertEqual(mock_first_client.await_count, 2)
725 # We should have triggered the retry logic at least once, but it
726 # might occur multiple times depending how long the first client
727 # takes to finish.
728 self.assertGreaterEqual(mock_second_client.await_count, 1)
730 # TODO DM-46204: This can be removed once the RSP recommended image has
731 # been upgraded to a version that contains DM-46129.
732 def test_deprecated_collection_endpoints(self):
733 # These REST endpoints are no longer used by Butler client so they need
734 # to be checked separately until they can be removed.
735 json = self.butler._connection.get(
736 "collection_info",
737 params={"name": "imported_g", "include_doc": True, "include_parents": True},
738 ).json()
739 self.assertEqual(json["name"], "imported_g")
740 self.assertEqual(json["type"], 1)
742 json = self.butler._connection.post(
743 "query_collections",
744 QueryCollectionsRequestModel(
745 search=["imported_*"], collection_types=[1], flatten_chains=False, include_chains=False
746 ),
747 ).json()
748 self.assertCountEqual(json["collections"], ["imported_g", "imported_r"])
750 def test_oversized_data_coordinate_upload(self):
751 with self.butler.query() as query:
752 ref = self.simple_dataset_ref
753 data_id = ref.dataId
754 data_coordinates = [DataCoordinate.standardize(data_id, visit=x) for x in range(100_001)]
755 with self.assertRaisesRegex(InvalidQueryError, "data coordinate rows"):
756 list(query.join_data_coordinates(data_coordinates).datasets(ref.datasetType, ref.run))
759@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
760class ButlerClientServerParityTestCase(unittest.TestCase):
761 """Test that `RemoteButler` and `DirectButler` agree for every combination
762 of component request and storage class override.
764 This uses its own server instance because
765 `ButlerClientServerTestCase` mocks the `DirectButler` datastore, which
766 would leave nothing to compare the `RemoteButler` results against.
767 """
769 COMPONENT = "summary"
771 COMPOSITE_STORAGE_CLASS = "MetricsConversion"
772 """Read storage class for the composite. Defines the same components as
773 ``StructuredData`` but with a different Python type.
774 """
776 COMPONENT_STORAGE_CLASS = "DictConvertibleModel"
777 """Read storage class for the component itself."""
779 DATA_ID = {"instrument": "DummyCamComp", "visit": 423}
780 COLLECTIONS = "ingest/run"
782 @classmethod
783 def setUpClass(cls):
784 server_instance = cls.enterClassContext(create_test_server(TESTDIR))
785 cls.direct_butler = server_instance.direct_butler
786 cls.remote_butler = server_instance.remote_butler
787 cls.storageClassFactory = StorageClassFactory()
788 repo = MetricTestRepo.create_from_butler(
789 server_instance.direct_butler, server_instance.config_file_path
790 )
791 # One dataset whose write storage class defines the component that will
792 # be requested, and one whose write storage class defines no components
793 # at all so that the component can only come from the read override.
794 cls.refs = {}
795 for name, storage_class in (
796 ("parity_with_components", "StructuredData"),
797 ("parity_no_components", "StructuredDataNoComponents"),
798 ):
799 dataset_type = addDatasetType(repo.butler, name, {"instrument", "visit"}, storage_class)
800 cls.refs[name] = repo.addDataset(cls.DATA_ID, datasetType=dataset_type)
802 def butlers(self):
803 """Iterate over the butlers that must behave identically.
805 Yields
806 ------
807 name : `str`
808 Name of the butler, for use in subtest labels.
809 butler : `Butler`
810 The butler itself.
811 """
812 yield "direct", self.direct_butler
813 yield "remote", self.remote_butler
815 def test_composite_parity(self):
816 """Both butlers must return the same composite, with and without a read
817 storage class override.
818 """
819 composite_pytype = self.storageClassFactory.getStorageClass(self.COMPOSITE_STORAGE_CLASS).pytype
820 for name, ref in self.refs.items():
821 write_pytype = ref.datasetType.storageClass.pytype
822 override_ref = ref.overrideStorageClass(self.COMPOSITE_STORAGE_CLASS)
823 cases = {
824 "get(ref)": (lambda b, r=ref: b.get(r), write_pytype),
825 "get(name, dataId)": (
826 lambda b, n=name: b.get(n, self.DATA_ID, collections=self.COLLECTIONS),
827 write_pytype,
828 ),
829 "get(ref, storageClass)": (
830 lambda b, r=ref: b.get(r, storageClass=self.COMPOSITE_STORAGE_CLASS),
831 composite_pytype,
832 ),
833 "get(override_ref)": (lambda b, r=override_ref: b.get(r), composite_pytype),
834 "getDeferred(ref).get()": (lambda b, r=ref: b.getDeferred(r).get(), write_pytype),
835 "getDeferred(ref, storageClass).get()": (
836 lambda b, r=ref: b.getDeferred(r, storageClass=self.COMPOSITE_STORAGE_CLASS).get(),
837 composite_pytype,
838 ),
839 }
840 for label, (call, pytype) in cases.items():
841 for butler_name, butler in self.butlers():
842 with self.subTest(dataset=name, case=label, butler=butler_name):
843 value = call(butler)
844 self.assertIsInstance(value, pytype)
845 self.assertEqual(value.summary, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
847 def test_component_parity(self):
848 """Both butlers must return the same component, whether it is defined
849 by the storage class used to write or only by the read override.
850 """
851 summary = MetricTestRepo.METRICS_EXAMPLE_SUMMARY
852 converted = DictConvertibleModel.from_dict(summary)
853 for name, ref in self.refs.items():
854 override_ref = ref.overrideStorageClass(self.COMPOSITE_STORAGE_CLASS)
855 component_ref = override_ref.makeComponentRef(self.COMPONENT)
856 cases = {
857 # A component of the overridden composite, requested by ref, by
858 # dataset type, and via a deferred handle. The component
859 # dataset type carries the composite override in its
860 # parentStorageClass, so unlike a component dataset type *name*
861 # it can express both overrides at once.
862 "get(component_ref)": (lambda b, r=component_ref: b.get(r), summary),
863 "get(component_ref, storageClass)": (
864 lambda b, r=component_ref: b.get(r, storageClass=self.COMPONENT_STORAGE_CLASS),
865 converted,
866 ),
867 "get(component_datasetType, dataId)": (
868 lambda b, r=component_ref: b.get(
869 r.datasetType, self.DATA_ID, collections=self.COLLECTIONS
870 ),
871 summary,
872 ),
873 "getDeferred(component_ref).get()": (
874 lambda b, r=component_ref: b.getDeferred(r).get(),
875 summary,
876 ),
877 # The handle storage class applies to the composite and so
878 # selects the component, while the one given to get() applies
879 # to the component itself.
880 "getDeferred(ref, storageClass).get(component)": (
881 lambda b, r=ref: b.getDeferred(r, storageClass=self.COMPOSITE_STORAGE_CLASS).get(
882 component=self.COMPONENT
883 ),
884 summary,
885 ),
886 "getDeferred(ref, storageClass).get(component, storageClass)": (
887 lambda b, r=ref: b.getDeferred(r, storageClass=self.COMPOSITE_STORAGE_CLASS).get(
888 component=self.COMPONENT, storageClass=self.COMPONENT_STORAGE_CLASS
889 ),
890 converted,
891 ),
892 }
893 if self.COMPONENT in ref.datasetType.storageClass.allComponents():
894 # The write storage class knows the component, so the plain
895 # forms that carry no override also work.
896 cases["get('name.component', dataId)"] = (
897 lambda b, n=name: b.get(
898 f"{n}.{self.COMPONENT}", self.DATA_ID, collections=self.COLLECTIONS
899 ),
900 summary,
901 )
902 cases["get('name.component', dataId, storageClass)"] = (
903 lambda b, n=name: b.get(
904 f"{n}.{self.COMPONENT}",
905 self.DATA_ID,
906 collections=self.COLLECTIONS,
907 storageClass=self.COMPONENT_STORAGE_CLASS,
908 ),
909 converted,
910 )
911 cases["getDeferred(ref).get(component)"] = (
912 lambda b, r=ref: b.getDeferred(r).get(component=self.COMPONENT),
913 summary,
914 )
915 cases["getDeferred('name.component', dataId).get()"] = (
916 lambda b, n=name: b.getDeferred(
917 f"{n}.{self.COMPONENT}", self.DATA_ID, collections=self.COLLECTIONS
918 ).get(),
919 summary,
920 )
921 for label, (call, expected) in cases.items():
922 for butler_name, butler in self.butlers():
923 with self.subTest(dataset=name, case=label, butler=butler_name):
924 self.assertEqual(call(butler), expected)
926 def test_unconveyable_component_override_parity(self):
927 """Both butlers must raise for a component that only the read storage
928 class defines when the request cannot convey the composite override.
930 A component dataset type *name* has nowhere to record a storage class
931 override for the composite, and the ``storageClass`` argument applies
932 to the component, so there is no way to say which composite storage
933 class defines the component. Raising is the only correct answer; a
934 `DatasetRef` or `DatasetType` must be used instead.
935 """
936 name = "parity_no_components"
937 ref = self.refs[name]
938 self.assertNotIn(self.COMPONENT, ref.datasetType.storageClass.allComponents())
939 component_type = f"{name}.{self.COMPONENT}"
940 cases = {
941 "get('name.component', dataId)": lambda b: b.get(
942 component_type, self.DATA_ID, collections=self.COLLECTIONS
943 ),
944 "get('name.component', dataId, storageClass)": lambda b: b.get(
945 component_type,
946 self.DATA_ID,
947 collections=self.COLLECTIONS,
948 storageClass=self.COMPONENT_STORAGE_CLASS,
949 ),
950 "getDeferred('name.component', dataId)": lambda b: b.getDeferred(
951 component_type, self.DATA_ID, collections=self.COLLECTIONS
952 ),
953 "getDeferred('name.component', dataId, storageClass)": lambda b: b.getDeferred(
954 component_type,
955 self.DATA_ID,
956 collections=self.COLLECTIONS,
957 storageClass=self.COMPONENT_STORAGE_CLASS,
958 ),
959 "getDeferred(ref).get(component)": lambda b: b.getDeferred(ref).get(component=self.COMPONENT),
960 }
961 for label, call in cases.items():
962 for butler_name, butler in self.butlers():
963 with self.subTest(case=label, butler=butler_name):
964 with self.assertRaises(UnknownComponentError):
965 call(butler)
968@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
969class ButlerClientServerAuthorizationTestCase(unittest.TestCase):
970 """Test authentication/authorization functionality."""
972 def test_group_authorization(self):
973 """Test that group membership repository authorization is checked when
974 repository is accessed.
975 """
976 with create_test_server(TESTDIR) as server_instance:
977 mock = MockGafaelfawrGroupAuthorizer()
978 server_instance.app.dependency_overrides[authorizer_dependency] = lambda: mock
979 server_instance.direct_butler.registry.registerDatasetType(
980 DatasetType("bias", [], "int", universe=server_instance.direct_butler.dimensions)
981 )
982 server_instance.direct_butler.collections.register("collection")
983 butler = server_instance.remote_butler
984 mock.set_response(False)
985 with self.assertRaises(ButlerServerError) as e:
986 butler.get_dataset_type("bias")
987 self.assertEqual(e.exception.status_code, 403)
988 with self.assertRaises(ButlerServerError) as e:
989 butler.query_datasets("bias", collections="*", find_first=False)
990 self.assertEqual(e.exception.status_code, 403)
992 mock.set_response(True)
993 self.assertEqual(butler.get_dataset_type("bias").name, "bias")
994 self.assertEqual(butler.query_datasets("bias", collections="collection", explain=False), [])
996 def test_cadc_auth(self) -> None:
997 """Test server running in CADC auth mode."""
998 with mock_config() as config:
999 config.authentication = "cadc"
1000 config.gafaelfawr_url = "DISABLED"
1001 with create_test_server(TESTDIR, server_config=config) as instance:
1002 self.assertIsInstance(instance.remote_butler._connection.auth, CadcAuthenticationProvider)
1004 # Set up a dataset backed by an HTTP URL.
1005 # CADC uses a plain HTTP service, not S3, for hosting Butler
1006 # artifacts.
1007 dataset_type = DatasetType("test", [], "int", universe=instance.direct_butler.dimensions)
1008 ref = DatasetRef(
1009 datasetType=dataset_type,
1010 dataId=DataCoordinate.makeEmpty(instance.direct_butler.dimensions),
1011 run="ingest/run",
1012 )
1013 path = ResourcePath("https://fake-server.example/some-directory/file.json")
1014 dataset = FileDataset(path, ref)
1015 # ingest() insists on doing file existence checks, and we don't
1016 # have an HTTP server to point it at.
1017 with unittest.mock.patch(
1018 "lsst.daf.butler.datastores.fileDatastore.FileDatastore._standardizeIngestPath"
1019 ) as mock:
1020 mock.return_value = path
1021 instance.direct_butler.ingest(dataset, transfer="direct", record_validation_info=False)
1023 # At the CADC, paths used for file download should NOT be a
1024 # signed URL, and should have authentication headers attached.
1025 def check_path(path_to_check: ResourcePath):
1026 self.assertEqual(str(path_to_check), str(path))
1027 assert isinstance(path_to_check, HttpResourcePath)
1028 self.assertIsNotNone(path_to_check._extra_headers)
1029 self.assertIsNotNone(path_to_check._extra_headers.get("Authorization"))
1031 check_path(instance.remote_butler.getURI(ref))
1032 transfer_map = instance.remote_butler._file_transfer_source.get_file_info_for_transfer(
1033 [ref.id]
1034 )
1035 check_path(transfer_map[ref.id][0].location.pathInStore)
1038def _create_corrupted_dataset(repo: MetricTestRepo) -> DatasetRef:
1039 run = "corrupted-run"
1040 ref = repo.addDataset({"instrument": "DummyCamComp", "visit": 423}, run=run)
1041 uris = repo.butler.getURIs(ref)
1042 oneOfTheComponents = list(uris.componentURIs.values())[0]
1043 oneOfTheComponents.write("corrupted data")
1044 return ref
1047def _create_simple_dataset(butler: Butler) -> DatasetRef:
1048 dataset_type = addDatasetType(butler, "test_int", {"instrument", "visit"}, "int")
1049 ref = butler.put(123, dataset_type, dataId={"instrument": "DummyCamComp", "visit": 423}, run="ingest/run")
1050 return ref
1053def _timeout_twice():
1054 """Return a mock side-effect function that raises a timeout error the first
1055 two times it is called.
1056 """
1057 count = 0
1059 def timeout(*args):
1060 nonlocal count
1061 count += 1
1062 if count <= 2:
1063 raise TimeoutError()
1064 return DEFAULT
1066 return timeout
1069class _MockStreamingQuery:
1070 @contextmanager
1071 def setup(self) -> Generator[None]:
1072 self.cleanup_executed = False
1073 self.count = 0
1074 try:
1075 yield
1076 finally:
1077 self.cleanup_executed = True
1079 def execute(self, context: None) -> Iterator[QueryKeepAliveModel]:
1080 for _ in range(10):
1081 self.count += 1
1082 yield QueryKeepAliveModel()
1085@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
1086class QueryStreamingTestCase(unittest.IsolatedAsyncioTestCase):
1087 """Test implementation details of query streaming code."""
1089 async def test_query_disconnects(self):
1090 query = _MockStreamingQuery()
1091 iter = _stream_query_pages(query, None)
1092 await anext(iter)
1093 await anext(iter)
1094 # FastAPI calls aclose() on the generator when disconnect occurs, so we
1095 # simulate that here.
1096 await iter.aclose()
1097 # We should have gone through the loop at least twice to yield the two
1098 # results we read. Can go up to 3-4 times depending on the timing of
1099 # the buffered results from the synchronous inner thread to the async
1100 # code.
1101 self.assertIn(query.count, (2, 3, 4))
1102 # Cleanup was run correctly, meaning the thread was given time to exit
1103 # before returning from aclose().
1104 self.assertTrue(query.cleanup_executed)
1107@unittest.skipIf(not butler_server_is_available, butler_server_import_error)
1108class QueryLimitsTestCase(unittest.IsolatedAsyncioTestCase):
1109 """Test details of the code that limits the maximum number of concurrent
1110 queries in the server.
1111 """
1113 async def test_query_limits(self):
1114 limits = lsst.daf.butler.remote_butler.server.handlers._query_limits.QueryLimits()
1116 await limits.enforce_query_limits("user1") # under limit, doesn't raise
1117 async with limits.track_query("user1"):
1118 await limits.enforce_query_limits("user1") # under limit, doesn't raise
1119 async with limits.track_query("user1"):
1120 with self.assertRaises(fastapi.HTTPException) as exc:
1121 await limits.enforce_query_limits("user1")
1122 self.assertEqual(exc.exception.status_code, 429)
1125if __name__ == "__main__":
1126 unittest.main()