Coverage for tests / test_server.py: 14%
365 statements
« prev ^ index » next coverage.py v7.13.2, created at 2026-02-08 16:37 -0800
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/>.
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
38from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel
40try:
41 # Failing to import any of these should disable the tests.
42 import fastapi
43 import httpx
44 import safir.dependencies.logger
45 from fastapi.testclient import TestClient
47 import lsst.daf.butler.remote_butler._query_results
48 import lsst.daf.butler.remote_butler.server.handlers._query_limits
49 import lsst.daf.butler.remote_butler.server.handlers._query_streaming
50 from lsst.daf.butler.remote_butler import ButlerServerError, RemoteButler
51 from lsst.daf.butler.remote_butler.authentication.cadc import CadcAuthenticationProvider
52 from lsst.daf.butler.remote_butler.authentication.rubin import (
53 _EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY,
54 RubinAuthenticationProvider,
55 )
56 from lsst.daf.butler.remote_butler.server import create_app
57 from lsst.daf.butler.remote_butler.server._config import mock_config
58 from lsst.daf.butler.remote_butler.server._dependencies import (
59 authorizer_dependency,
60 butler_factory_dependency,
61 )
62 from lsst.daf.butler.remote_butler.server._gafaelfawr import MockGafaelfawrGroupAuthorizer
63 from lsst.daf.butler.remote_butler.server.handlers._utils import generate_file_download_uri
64 from lsst.daf.butler.remote_butler.server_models import QueryCollectionsRequestModel
65 from lsst.daf.butler.tests.server import TEST_REPOSITORY_NAME, UnhandledServerError, create_test_server
67 reason_text = ""
68except ImportError as e:
69 create_test_server = None
70 reason_text = str(e)
73from lsst.daf.butler import (
74 Butler,
75 DataCoordinate,
76 DatasetId,
77 DatasetNotFoundError,
78 DatasetRef,
79 DatasetType,
80 FileDataset,
81 InvalidQueryError,
82 LabeledButlerFactory,
83 MissingDatasetTypeError,
84 NoDefaultCollectionError,
85 StorageClassFactory,
86)
87from lsst.daf.butler.datastore import DatasetRefURIs
88from lsst.daf.butler.registry import RegistryDefaults
89from lsst.daf.butler.tests import DatastoreMock, addDatasetType
90from lsst.daf.butler.tests.utils import MetricsExample, MetricTestRepo, mock_env
91from lsst.resources import ResourcePath
92from lsst.resources.http import HttpResourcePath
94TESTDIR = os.path.abspath(os.path.dirname(__file__))
97@unittest.skipIf(create_test_server is None, f"Server dependencies not installed: {reason_text}")
98class ButlerClientServerTestCase(unittest.TestCase):
99 """Test for Butler client/server."""
101 @classmethod
102 def setUpClass(cls):
103 server_instance = cls.enterClassContext(create_test_server(TESTDIR))
104 cls.server_instance = server_instance
105 cls.client = server_instance.client
106 cls.butler = server_instance.remote_butler
107 cls.butler_without_error_propagation = server_instance.remote_butler_without_error_propagation
109 cls.storageClassFactory = StorageClassFactory()
111 cls.repo = MetricTestRepo.create_from_butler(
112 server_instance.direct_butler, server_instance.config_file_path
113 )
114 # Add a file with corrupted data for testing error conditions
115 cls.dataset_with_corrupted_data = _create_corrupted_dataset(cls.repo)
116 # All of the datasets that come with MetricTestRepo are disassembled
117 # composites. Add a simple dataset for testing the common case.
118 cls.simple_dataset_ref = _create_simple_dataset(server_instance.direct_butler)
120 # Populate the test server.
121 # The DatastoreMock is required because the datasets referenced in
122 # these imports do not point at real files.
123 direct_butler = server_instance.direct_butler
124 DatastoreMock.apply(direct_butler)
125 direct_butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml")
126 direct_butler.import_(filename="resource://lsst.daf.butler/tests/registry_data/datasets.yaml")
128 def test_health_check(self):
129 try:
130 import importlib.metadata
132 importlib.metadata.metadata("lsst.daf.butler")
133 except ModuleNotFoundError:
134 raise self.skipTest("Standard python package metadata not available. Butler not pip installed.")
135 response = self.client.get("/")
136 self.assertEqual(response.status_code, 200)
137 self.assertEqual(response.json()["name"], "butler")
139 def test_static_files(self):
140 with tempfile.TemporaryDirectory() as tmpdir:
141 with open(os.path.join(tmpdir, "temp.txt"), "w") as fh:
142 fh.write("test data 123")
144 with mock_config() as server_config:
145 server_config.static_files_path = tmpdir
146 with create_test_server(TESTDIR, server_config=server_config) as server:
147 response = server.client.get("/api/butler/configs/temp.txt")
148 self.assertEqual(response.status_code, 200)
149 self.assertEqual(response.text, "test data 123")
151 def test_dimension_universe(self):
152 universe = self.butler.dimensions
153 self.assertEqual(universe.namespace, "daf_butler")
155 def test_get_dataset_type(self):
156 bias_type = self.butler.get_dataset_type("bias")
157 self.assertEqual(bias_type.name, "bias")
159 with self.assertRaises(MissingDatasetTypeError):
160 self.butler_without_error_propagation.get_dataset_type("not_bias")
162 def test_find_dataset(self):
163 storage_class = self.storageClassFactory.getStorageClass("Exposure")
165 ref = self.butler.find_dataset("bias", collections="imported_g", detector=1, instrument="Cam1")
166 self.assertIsInstance(ref, DatasetRef)
167 self.assertEqual(ref.id, uuid.UUID("e15ab039-bc8b-4135-87c5-90902a7c0b22"))
168 self.assertFalse(ref.dataId.hasRecords())
170 # Try again with variation of parameters.
171 ref_new = self.butler.find_dataset(
172 "bias",
173 {"detector": 1},
174 collections="imported_g",
175 instrument="Cam1",
176 dimension_records=True,
177 )
178 self.assertEqual(ref_new, ref)
179 self.assertTrue(ref_new.dataId.hasRecords())
181 ref_new = self.butler.find_dataset(
182 ref.datasetType,
183 DataCoordinate.standardize(detector=1, instrument="Cam1", universe=self.butler.dimensions),
184 collections="imported_g",
185 storage_class=storage_class,
186 )
187 self.assertEqual(ref_new, ref)
189 ref2 = self.butler.get_dataset(ref.id)
190 self.assertEqual(ref2, ref)
192 # Use detector name to find it.
193 ref3 = self.butler.find_dataset(
194 ref.datasetType,
195 collections="imported_g",
196 instrument="Cam1",
197 full_name="Aa",
198 )
199 self.assertEqual(ref2, ref3)
201 # Try expanded refs.
202 self.assertFalse(ref.dataId.hasRecords())
203 expanded = self.butler.get_dataset(ref.id, dimension_records=True)
204 self.assertTrue(expanded.dataId.hasRecords())
206 # The test datasets are all Exposure so storage class conversion
207 # can not be tested until we fix that. For now at least test the
208 # code paths.
209 bias = self.butler.get_dataset(ref.id, storage_class=storage_class)
210 self.assertEqual(bias.datasetType.storageClass, storage_class)
212 # Unknown dataset should not fail.
213 self.assertIsNone(self.butler.get_dataset(uuid.uuid4()))
214 self.assertIsNone(self.butler.get_dataset(uuid.uuid4(), storage_class="NumpyArray"))
216 def test_instantiate_via_butler_http_search(self):
217 """Ensure that the primary Butler constructor's automatic search logic
218 correctly locates and reads the configuration file and ends up with a
219 RemoteButler pointing to the correct URL
220 """
222 # This is kind of a fragile test. Butler's search logic does a lot of
223 # manipulations involving creating new ResourcePaths, and ResourcePath
224 # doesn't use httpx so we can't easily inject the TestClient in there.
225 # We don't have an actual valid HTTP URL to give to the constructor
226 # because the test instance of the server is accessed via ASGI.
227 #
228 # Instead we just monkeypatch the HTTPResourcePath 'read' method and
229 # hope that all ResourcePath HTTP reads during construction are going
230 # to the server under test.
231 def override_read(http_resource_path):
232 return self.client.get(http_resource_path.geturl()).content
234 server_url = f"https://test.example/api/butler/repo/{TEST_REPOSITORY_NAME}/"
236 with patch.object(HttpResourcePath, "read", override_read):
237 # RegistryDefaults.finish() needs to download the dimension
238 # universe from the server, which will fail because there is no
239 # server here. So mock it out.
240 with patch.object(RegistryDefaults, "finish"):
241 # Add access key to environment variables. RemoteButler
242 # instantiation will throw an error if access key is not
243 # available.
244 with mock_env({_EXPLICIT_BUTLER_ACCESS_TOKEN_ENVIRONMENT_KEY: "fake-access-token"}):
245 butler = Butler(
246 server_url,
247 collections=["collection1", "collection2"],
248 run="collection2",
249 )
250 self.enterContext(butler)
251 self.assertIsInstance(butler, RemoteButler)
252 self.assertEqual(butler._connection.server_url, server_url)
253 self.assertEqual(butler.collections.defaults, ("collection1", "collection2"))
254 self.assertEqual(butler.run, "collection2")
255 # A butler created this way uses the default cache config.
256 self.assertFalse(butler._use_disabled_datastore_cache)
258 butler_factory = LabeledButlerFactory({"server": server_url})
259 factory_created_butler = butler_factory.create_butler(label="server", access_token="token")
260 self.assertIsInstance(factory_created_butler, RemoteButler)
261 self.assertTrue(factory_created_butler._use_disabled_datastore_cache)
262 self.assertEqual(factory_created_butler._connection.server_url, server_url)
264 def test_get(self):
265 dataset_type = "test_metric_comp"
266 data_id = {"instrument": "DummyCamComp", "visit": 423}
267 collections = "ingest/run"
268 # Test get() of a DatasetRef.
269 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
270 metric = self.butler.get(ref)
271 self.assertIsInstance(metric, MetricsExample)
272 self.assertEqual(metric.summary, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
274 # Test get() by DataId.
275 data_id_metric = self.butler.get(dataset_type, dataId=data_id, collections=collections)
276 self.assertEqual(metric, data_id_metric)
277 # Test get() by DataId dict augmented with kwargs.
278 kwarg_metric = self.butler.get(
279 dataset_type, dataId={"instrument": "DummyCamComp"}, collections=collections, visit=423
280 )
281 self.assertEqual(metric, kwarg_metric)
282 # Test get() by DataId DataCoordinate augmented with kwargs.
283 coordinate = DataCoordinate.make_empty(self.butler.dimensions)
284 kwarg_data_coordinate_metric = self.butler.get(
285 dataset_type, dataId=coordinate, collections=collections, instrument="DummyCamComp", visit=423
286 )
287 self.assertEqual(metric, kwarg_data_coordinate_metric)
288 # Test get() of a non-existent DataId.
289 invalid_data_id = {"instrument": "NotAValidlInstrument", "visit": 423}
290 with self.assertRaises(DatasetNotFoundError):
291 self.butler_without_error_propagation.get(
292 dataset_type, dataId=invalid_data_id, collections=collections
293 )
295 # Test get() by DataId with default collections.
296 butler_with_default_collection = self.butler.clone(collections="ingest/run")
297 default_collection_metric = butler_with_default_collection.get(dataset_type, dataId=data_id)
298 self.assertEqual(metric, default_collection_metric)
300 # Test get() by DataId with no collections specified.
301 with self.assertRaises(NoDefaultCollectionError):
302 self.butler_without_error_propagation.get(dataset_type, dataId=data_id)
304 # Test looking up a non-existent ref
305 invalid_ref = ref.replace(id=uuid.uuid4())
306 with self.assertRaises(DatasetNotFoundError):
307 self.butler_without_error_propagation.get(invalid_ref)
309 with self.assertRaises(RuntimeError):
310 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data)
312 # Test storage class override
313 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
315 def check_sc_override(converted):
316 self.assertNotEqual(type(metric), type(converted))
317 self.assertIsInstance(converted, new_sc.pytype)
318 self.assertEqual(metric, converted)
320 check_sc_override(self.butler.get(ref, storageClass=new_sc))
322 # Test storage class override via DatasetRef.
323 check_sc_override(self.butler.get(ref.overrideStorageClass("MetricsConversion")))
324 # Test storage class override via DatasetType.
325 check_sc_override(
326 self.butler.get(
327 ref.datasetType.overrideStorageClass(new_sc), dataId=data_id, collections=collections
328 )
329 )
331 # Test component override via DatasetRef.
332 component_ref = ref.makeComponentRef("summary")
333 component_data = self.butler.get(component_ref)
334 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
336 # Test overriding both storage class and component via DatasetRef.
337 converted_component_data = self.butler.get(component_ref, storageClass="DictConvertibleModel")
338 self.assertIsInstance(converted_component_data, DictConvertibleModel)
339 self.assertEqual(converted_component_data.content, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
341 # Test component override via DatasetType.
342 dataset_type_component_data = self.butler.get(
343 component_ref.datasetType, component_ref.dataId, collections=collections
344 )
345 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
347 def test_getURIs_no_components(self):
348 # This dataset does not have components, and should return one URI.
349 def check_uri(uri: ResourcePath):
350 self.assertIsNotNone(uris.primaryURI)
351 self.assertEqual(uris.primaryURI.scheme, "https")
352 self.assertEqual(uris.primaryURI.read(), b"123")
354 uris = self.butler.getURIs(self.simple_dataset_ref)
355 self.assertEqual(len(uris.componentURIs), 0)
356 check_uri(uris.primaryURI)
358 check_uri(self.butler.getURI(self.simple_dataset_ref))
360 def test_getURIs_multiple_components(self):
361 # This dataset has multiple components, so we should get back multiple
362 # URIs.
363 dataset_type = "test_metric_comp"
364 data_id = {"instrument": "DummyCamComp", "visit": 423}
365 collections = "ingest/run"
367 def check_uris(uris: DatasetRefURIs):
368 self.assertIsNone(uris.primaryURI)
369 self.assertEqual(len(uris.componentURIs), 3)
370 path = uris.componentURIs["summary"]
371 self.assertEqual(path.scheme, "https")
372 data = path.read()
373 self.assertEqual(data, b"AM1: 5.2\nAM2: 30.6\n")
375 uris = self.butler.getURIs(dataset_type, dataId=data_id, collections=collections)
376 check_uris(uris)
378 # Calling getURI on a multi-file dataset raises an exception
379 with self.assertRaises(RuntimeError):
380 self.butler.getURI(dataset_type, dataId=data_id, collections=collections)
382 # getURIs does NOT respect component overrides on the DatasetRef,
383 # instead returning the parent's URIs. Unclear if this is "correct"
384 # from a conceptual point of view, but this matches DirectButler
385 # behavior.
386 ref = self.butler.find_dataset(dataset_type, data_id=data_id, collections=collections)
387 componentRef = ref.makeComponentRef("summary")
388 componentUris = self.butler.getURIs(componentRef)
389 check_uris(componentUris)
391 def test_file_download_redirect(self):
392 def get_download_redirect(id: DatasetId, component: str | None = None) -> httpx.Response:
393 uri = generate_file_download_uri("http://unittest.test/", TEST_REPOSITORY_NAME, id, component)
394 return self.client.get(
395 uri,
396 follow_redirects=False,
397 headers=RubinAuthenticationProvider("mock-token").get_server_headers(),
398 )
400 # Test behavior of a single-file dataset.
401 response = get_download_redirect(self.simple_dataset_ref.id)
402 self.assertEqual(response.status_code, 307)
403 self.assertTrue(response.has_redirect_location)
404 assert response.next_request is not None
405 self.assertEqual(response.next_request.url.scheme, "https")
406 self.assertIn("test_int_DummyCamComp_R_d-r_423_ingest_run.json", response.next_request.url.path)
408 response = get_download_redirect(self.simple_dataset_ref.id, "somecomponent")
409 self.assertEqual(response.status_code, 404)
411 # This dataset is a "disassembled composite" with multiple files.
412 dataset_type = "test_metric_comp"
413 data_id = {"instrument": "DummyCamComp", "visit": 423}
414 collections = "ingest/run"
415 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
417 # Getting single component of a multi-file "disassembled composite".
418 response = get_download_redirect(ref.id, "summary")
419 self.assertEqual(response.status_code, 307)
420 self.assertTrue(response.has_redirect_location)
421 assert response.next_request is not None
422 self.assertEqual(response.next_request.url.scheme, "https")
423 self.assertIn("test_metric_comp.summary", response.next_request.url.path)
425 # Unknown component.
426 response = get_download_redirect(ref.id, "badcomponent")
427 self.assertEqual(response.status_code, 404)
429 # Not specifying the component for a multi-file "disassembled
430 # composite".
431 response = get_download_redirect(ref.id, None)
432 self.assertEqual(response.status_code, 422)
434 # Unknown dataset.
435 response = get_download_redirect(uuid.UUID("59467c1b-fa13-4f7a-8ff8-cd83e092e563"))
436 self.assertEqual(response.status_code, 404)
438 def test_auth_check(self):
439 # This is checking that the unit-test middleware for validating the
440 # authentication headers is working. It doesn't test actual server
441 # functionality -- in a real deployment, the authentication headers are
442 # handled by GafaelfawrIngress, not our app.
443 with self.assertRaises(UnhandledServerError):
444 self.client.get("/v1/dataset_type/int")
446 def test_exception_logging(self):
447 app = create_app()
449 def raise_error():
450 raise RuntimeError("An unhandled error")
452 app.dependency_overrides[butler_factory_dependency] = raise_error
453 client = TestClient(app, raise_server_exceptions=False)
455 with patch.object(safir.dependencies.logger, "logger_dependency") as mock_logger_dep:
456 mock_logger = NonCallableMock(["aerror"])
458 async def noop():
459 pass
461 mock_logger.aerror.return_value = noop()
463 async def get_logger():
464 return mock_logger
466 mock_logger_dep.return_value = get_logger()
467 client.get(
468 "/api/butler/repo/something/v1/dataset_type/int",
469 headers={"X-Auth-Request-User": "user-name", "X-Butler-Client-Request-Id": "request-id"},
470 )
471 mock_logger_dep.assert_called_once()
473 mock_logger.aerror.assert_called_once()
474 args, kwargs = mock_logger.aerror.call_args
475 self.assertIsInstance(kwargs["exc_info"], RuntimeError)
476 self.assertEqual(kwargs["clientRequestId"], "request-id")
477 self.assertEqual(kwargs["user"], "user-name")
479 def test_query_keepalive(self):
480 """Test that long-running queries stream keep-alive messages to stop
481 the HTTP connection from closing before they are able to return
482 results.
483 """
484 # Normally it takes 15 seconds for a timeout -- mock it to trigger
485 # immediately instead.
486 with patch.object(
487 lsst.daf.butler.remote_butler.server.handlers._query_streaming, "_timeout"
488 ) as mock_timeout:
489 # Hook into QueryDriver to track the number of keep-alives we have
490 # seen.
491 with patch.object(
492 lsst.daf.butler.remote_butler._query_results, "_received_keep_alive"
493 ) as mock_keep_alive:
494 mock_timeout.side_effect = _timeout_twice()
495 with self.butler.query() as query:
496 datasets = list(query.datasets("bias", "imported_g"))
497 self.assertEqual(len(datasets), 3)
498 self.assertGreaterEqual(mock_timeout.call_count, 3)
499 self.assertGreaterEqual(mock_keep_alive.call_count, 2)
501 def test_query_retries(self):
502 """Test that the server will send HTTP status 503 to put backpressure
503 on clients if it is overloaded, and that the client will retry if this
504 happens.
505 """
506 query_event = threading.Event()
507 retry_event = asyncio.Event()
509 async def block_first_request() -> None:
510 # Signal the unit tests that we have reached the critical section
511 # in the server, where the first client has reserved the query
512 # slot.
513 query_event.set()
514 # Block inside the query, until the 2nd client has been forced to
515 # retry.
516 await retry_event.wait()
518 async def block_second_request() -> None:
519 # Release the first client, so it can finish its query and prevent
520 # this client from being blocked on the next go-round.
521 retry_event.set()
523 def do_query(butler: Butler) -> list[DatasetRef]:
524 return butler.query_datasets("bias", "imported_g")
526 with (
527 patch.object(
528 lsst.daf.butler.remote_butler.server.handlers._query_limits,
529 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES",
530 new=1,
531 ),
532 patch.object(
533 lsst.daf.butler.remote_butler.server.handlers._query_limits, "_QUERY_RETRY_SECONDS", new=1
534 ),
535 patch.object(
536 lsst.daf.butler.remote_butler.server.handlers._query_limits,
537 "_block_query_for_unit_test",
538 new=AsyncMock(wraps=block_first_request),
539 ) as mock_first_client,
540 patch.object(
541 lsst.daf.butler.remote_butler.server.handlers._query_limits,
542 "_block_retry_for_unit_test",
543 new=AsyncMock(wraps=block_second_request),
544 ) as mock_second_client,
545 ThreadPoolExecutor(max_workers=1) as exec1,
546 ThreadPoolExecutor(max_workers=1) as exec2,
547 ):
548 first_butler = self.butler
549 second_butler = self.butler.clone()
551 # Run the first client up until the server starts executing its
552 # query.
553 future1 = exec1.submit(do_query, first_butler)
554 event_reached = query_event.wait(60)
555 if not event_reached:
556 raise TimeoutError("Server did not execute query logic as expected.")
558 # Start the second client, which will trigger the retry logic and
559 # release the first client to finish its query.
560 future2 = exec2.submit(do_query, second_butler)
562 result1 = future1.result(60)
563 result2 = future2.result(60)
564 self.assertEqual(len(result1), 3)
565 self.assertEqual(len(result2), 3)
566 # The original thread should have gone through this section, and
567 # then the 2nd thread after it retries.
568 self.assertEqual(mock_first_client.await_count, 2)
569 # We should have triggered the retry logic at least once, but it
570 # might occur multiple times depending how long the first client
571 # takes to finish.
572 self.assertGreaterEqual(mock_second_client.await_count, 1)
574 # TODO DM-46204: This can be removed once the RSP recommended image has
575 # been upgraded to a version that contains DM-46129.
576 def test_deprecated_collection_endpoints(self):
577 # These REST endpoints are no longer used by Butler client so they need
578 # to be checked separately until they can be removed.
579 json = self.butler._connection.get(
580 "collection_info",
581 params={"name": "imported_g", "include_doc": True, "include_parents": True},
582 ).json()
583 self.assertEqual(json["name"], "imported_g")
584 self.assertEqual(json["type"], 1)
586 json = self.butler._connection.post(
587 "query_collections",
588 QueryCollectionsRequestModel(
589 search=["imported_*"], collection_types=[1], flatten_chains=False, include_chains=False
590 ),
591 ).json()
592 self.assertCountEqual(json["collections"], ["imported_g", "imported_r"])
594 def test_oversized_data_coordinate_upload(self):
595 with self.butler.query() as query:
596 ref = self.simple_dataset_ref
597 data_id = ref.dataId
598 data_coordinates = [DataCoordinate.standardize(data_id, visit=x) for x in range(100_001)]
599 with self.assertRaisesRegex(InvalidQueryError, "data coordinate rows"):
600 list(query.join_data_coordinates(data_coordinates).datasets(ref.datasetType, ref.run))
603@unittest.skipIf(create_test_server is None, f"Server dependencies not installed: {reason_text}")
604class ButlerClientServerAuthorizationTestCase(unittest.TestCase):
605 """Test authentication/authorization functionality."""
607 def test_group_authorization(self):
608 """Test that group membership repository authorization is checked when
609 repository is accessed.
610 """
611 with create_test_server(TESTDIR) as server_instance:
612 mock = MockGafaelfawrGroupAuthorizer()
613 server_instance.app.dependency_overrides[authorizer_dependency] = lambda: mock
614 server_instance.direct_butler.registry.registerDatasetType(
615 DatasetType("bias", [], "int", universe=server_instance.direct_butler.dimensions)
616 )
617 server_instance.direct_butler.collections.register("collection")
618 butler = server_instance.remote_butler
619 mock.set_response(False)
620 with self.assertRaises(ButlerServerError) as e:
621 butler.get_dataset_type("bias")
622 self.assertEqual(e.exception.status_code, 403)
623 with self.assertRaises(ButlerServerError) as e:
624 butler.query_datasets("bias", collections="*", find_first=False)
625 self.assertEqual(e.exception.status_code, 403)
627 mock.set_response(True)
628 self.assertEqual(butler.get_dataset_type("bias").name, "bias")
629 self.assertEqual(butler.query_datasets("bias", collections="collection", explain=False), [])
631 def test_cadc_auth(self) -> None:
632 """Test server running in CADC auth mode."""
633 with mock_config() as config:
634 config.authentication = "cadc"
635 config.gafaelfawr_url = "DISABLED"
636 with create_test_server(TESTDIR, server_config=config) as instance:
637 self.assertIsInstance(instance.remote_butler._connection.auth, CadcAuthenticationProvider)
639 # Set up a dataset backed by an HTTP URL.
640 # CADC uses a plain HTTP service, not S3, for hosting Butler
641 # artifacts.
642 dataset_type = DatasetType("test", [], "int", universe=instance.direct_butler.dimensions)
643 ref = DatasetRef(
644 datasetType=dataset_type,
645 dataId=DataCoordinate.makeEmpty(instance.direct_butler.dimensions),
646 run="ingest/run",
647 )
648 path = ResourcePath("https://fake-server.example/some-directory/file.json")
649 dataset = FileDataset(path, ref)
650 # ingest() insists on doing file existence checks, and we don't
651 # have an HTTP server to point it at.
652 with unittest.mock.patch(
653 "lsst.daf.butler.datastores.fileDatastore.FileDatastore._standardizeIngestPath"
654 ) as mock:
655 mock.return_value = path
656 instance.direct_butler.ingest(dataset, transfer="direct", record_validation_info=False)
658 # At the CADC, paths used for file download should NOT be a
659 # signed URL, and should have authentication headers attached.
660 def check_path(path_to_check: ResourcePath):
661 self.assertEqual(str(path_to_check), str(path))
662 assert isinstance(path_to_check, HttpResourcePath)
663 self.assertIsNotNone(path_to_check._extra_headers)
664 self.assertIsNotNone(path_to_check._extra_headers.get("Authorization"))
666 check_path(instance.remote_butler.getURI(ref))
667 transfer_map = instance.remote_butler._file_transfer_source.get_file_info_for_transfer(
668 [ref.id]
669 )
670 check_path(transfer_map[ref.id][0].location.pathInStore)
673def _create_corrupted_dataset(repo: MetricTestRepo) -> DatasetRef:
674 run = "corrupted-run"
675 ref = repo.addDataset({"instrument": "DummyCamComp", "visit": 423}, run=run)
676 uris = repo.butler.getURIs(ref)
677 oneOfTheComponents = list(uris.componentURIs.values())[0]
678 oneOfTheComponents.write("corrupted data")
679 return ref
682def _create_simple_dataset(butler: Butler) -> DatasetRef:
683 dataset_type = addDatasetType(butler, "test_int", {"instrument", "visit"}, "int")
684 ref = butler.put(123, dataset_type, dataId={"instrument": "DummyCamComp", "visit": 423}, run="ingest/run")
685 return ref
688def _timeout_twice():
689 """Return a mock side-effect function that raises a timeout error the first
690 two times it is called.
691 """
692 count = 0
694 def timeout(*args):
695 nonlocal count
696 count += 1
697 if count <= 2:
698 raise TimeoutError()
699 return DEFAULT
701 return timeout
704@unittest.skipIf(create_test_server is None, f"Server dependencies not installed: {reason_text}")
705class QueryLimitsTestCase(unittest.IsolatedAsyncioTestCase):
706 """Test details of the code that limits the maximum number of concurrent
707 queries in the server.
708 """
710 async def test_query_limits(self):
711 limits = lsst.daf.butler.remote_butler.server.handlers._query_limits.QueryLimits()
713 await limits.enforce_query_limits("user1") # under limit, doesn't raise
714 async with limits.track_query("user1"):
715 await limits.enforce_query_limits("user1") # under limit, doesn't raise
716 async with limits.track_query("user1"):
717 with self.assertRaises(fastapi.HTTPException) as exc:
718 await limits.enforce_query_limits("user1")
719 self.assertEqual(exc.exception.status_code, 429)
722if __name__ == "__main__":
723 unittest.main()