Coverage for tests/test_server.py: 14%
364 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:22 -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/>.
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.assertIsInstance(butler, RemoteButler)
251 self.assertEqual(butler._connection.server_url, server_url)
252 self.assertEqual(butler.collections.defaults, ("collection1", "collection2"))
253 self.assertEqual(butler.run, "collection2")
254 # A butler created this way uses the default cache config.
255 self.assertFalse(butler._use_disabled_datastore_cache)
257 butler_factory = LabeledButlerFactory({"server": server_url})
258 factory_created_butler = butler_factory.create_butler(label="server", access_token="token")
259 self.assertIsInstance(factory_created_butler, RemoteButler)
260 self.assertTrue(factory_created_butler._use_disabled_datastore_cache)
261 self.assertEqual(factory_created_butler._connection.server_url, server_url)
263 def test_get(self):
264 dataset_type = "test_metric_comp"
265 data_id = {"instrument": "DummyCamComp", "visit": 423}
266 collections = "ingest/run"
267 # Test get() of a DatasetRef.
268 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
269 metric = self.butler.get(ref)
270 self.assertIsInstance(metric, MetricsExample)
271 self.assertEqual(metric.summary, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
273 # Test get() by DataId.
274 data_id_metric = self.butler.get(dataset_type, dataId=data_id, collections=collections)
275 self.assertEqual(metric, data_id_metric)
276 # Test get() by DataId dict augmented with kwargs.
277 kwarg_metric = self.butler.get(
278 dataset_type, dataId={"instrument": "DummyCamComp"}, collections=collections, visit=423
279 )
280 self.assertEqual(metric, kwarg_metric)
281 # Test get() by DataId DataCoordinate augmented with kwargs.
282 coordinate = DataCoordinate.make_empty(self.butler.dimensions)
283 kwarg_data_coordinate_metric = self.butler.get(
284 dataset_type, dataId=coordinate, collections=collections, instrument="DummyCamComp", visit=423
285 )
286 self.assertEqual(metric, kwarg_data_coordinate_metric)
287 # Test get() of a non-existent DataId.
288 invalid_data_id = {"instrument": "NotAValidlInstrument", "visit": 423}
289 with self.assertRaises(DatasetNotFoundError):
290 self.butler_without_error_propagation.get(
291 dataset_type, dataId=invalid_data_id, collections=collections
292 )
294 # Test get() by DataId with default collections.
295 butler_with_default_collection = self.butler.clone(collections="ingest/run")
296 default_collection_metric = butler_with_default_collection.get(dataset_type, dataId=data_id)
297 self.assertEqual(metric, default_collection_metric)
299 # Test get() by DataId with no collections specified.
300 with self.assertRaises(NoDefaultCollectionError):
301 self.butler_without_error_propagation.get(dataset_type, dataId=data_id)
303 # Test looking up a non-existent ref
304 invalid_ref = ref.replace(id=uuid.uuid4())
305 with self.assertRaises(DatasetNotFoundError):
306 self.butler_without_error_propagation.get(invalid_ref)
308 with self.assertRaises(RuntimeError):
309 self.butler_without_error_propagation.get(self.dataset_with_corrupted_data)
311 # Test storage class override
312 new_sc = self.storageClassFactory.getStorageClass("MetricsConversion")
314 def check_sc_override(converted):
315 self.assertNotEqual(type(metric), type(converted))
316 self.assertIsInstance(converted, new_sc.pytype)
317 self.assertEqual(metric, converted)
319 check_sc_override(self.butler.get(ref, storageClass=new_sc))
321 # Test storage class override via DatasetRef.
322 check_sc_override(self.butler.get(ref.overrideStorageClass("MetricsConversion")))
323 # Test storage class override via DatasetType.
324 check_sc_override(
325 self.butler.get(
326 ref.datasetType.overrideStorageClass(new_sc), dataId=data_id, collections=collections
327 )
328 )
330 # Test component override via DatasetRef.
331 component_ref = ref.makeComponentRef("summary")
332 component_data = self.butler.get(component_ref)
333 self.assertEqual(component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
335 # Test overriding both storage class and component via DatasetRef.
336 converted_component_data = self.butler.get(component_ref, storageClass="DictConvertibleModel")
337 self.assertIsInstance(converted_component_data, DictConvertibleModel)
338 self.assertEqual(converted_component_data.content, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
340 # Test component override via DatasetType.
341 dataset_type_component_data = self.butler.get(
342 component_ref.datasetType, component_ref.dataId, collections=collections
343 )
344 self.assertEqual(dataset_type_component_data, MetricTestRepo.METRICS_EXAMPLE_SUMMARY)
346 def test_getURIs_no_components(self):
347 # This dataset does not have components, and should return one URI.
348 def check_uri(uri: ResourcePath):
349 self.assertIsNotNone(uris.primaryURI)
350 self.assertEqual(uris.primaryURI.scheme, "https")
351 self.assertEqual(uris.primaryURI.read(), b"123")
353 uris = self.butler.getURIs(self.simple_dataset_ref)
354 self.assertEqual(len(uris.componentURIs), 0)
355 check_uri(uris.primaryURI)
357 check_uri(self.butler.getURI(self.simple_dataset_ref))
359 def test_getURIs_multiple_components(self):
360 # This dataset has multiple components, so we should get back multiple
361 # URIs.
362 dataset_type = "test_metric_comp"
363 data_id = {"instrument": "DummyCamComp", "visit": 423}
364 collections = "ingest/run"
366 def check_uris(uris: DatasetRefURIs):
367 self.assertIsNone(uris.primaryURI)
368 self.assertEqual(len(uris.componentURIs), 3)
369 path = uris.componentURIs["summary"]
370 self.assertEqual(path.scheme, "https")
371 data = path.read()
372 self.assertEqual(data, b"AM1: 5.2\nAM2: 30.6\n")
374 uris = self.butler.getURIs(dataset_type, dataId=data_id, collections=collections)
375 check_uris(uris)
377 # Calling getURI on a multi-file dataset raises an exception
378 with self.assertRaises(RuntimeError):
379 self.butler.getURI(dataset_type, dataId=data_id, collections=collections)
381 # getURIs does NOT respect component overrides on the DatasetRef,
382 # instead returning the parent's URIs. Unclear if this is "correct"
383 # from a conceptual point of view, but this matches DirectButler
384 # behavior.
385 ref = self.butler.find_dataset(dataset_type, data_id=data_id, collections=collections)
386 componentRef = ref.makeComponentRef("summary")
387 componentUris = self.butler.getURIs(componentRef)
388 check_uris(componentUris)
390 def test_file_download_redirect(self):
391 def get_download_redirect(id: DatasetId, component: str | None = None) -> httpx.Response:
392 uri = generate_file_download_uri("http://unittest.test/", TEST_REPOSITORY_NAME, id, component)
393 return self.client.get(
394 uri,
395 follow_redirects=False,
396 headers=RubinAuthenticationProvider("mock-token").get_server_headers(),
397 )
399 # Test behavior of a single-file dataset.
400 response = get_download_redirect(self.simple_dataset_ref.id)
401 self.assertEqual(response.status_code, 307)
402 self.assertTrue(response.has_redirect_location)
403 assert response.next_request is not None
404 self.assertEqual(response.next_request.url.scheme, "https")
405 self.assertIn("test_int_DummyCamComp_R_d-r_423_ingest_run.json", response.next_request.url.path)
407 response = get_download_redirect(self.simple_dataset_ref.id, "somecomponent")
408 self.assertEqual(response.status_code, 404)
410 # This dataset is a "disassembled composite" with multiple files.
411 dataset_type = "test_metric_comp"
412 data_id = {"instrument": "DummyCamComp", "visit": 423}
413 collections = "ingest/run"
414 ref = self.butler.find_dataset(dataset_type, data_id, collections=collections)
416 # Getting single component of a multi-file "disassembled composite".
417 response = get_download_redirect(ref.id, "summary")
418 self.assertEqual(response.status_code, 307)
419 self.assertTrue(response.has_redirect_location)
420 assert response.next_request is not None
421 self.assertEqual(response.next_request.url.scheme, "https")
422 self.assertIn("test_metric_comp.summary", response.next_request.url.path)
424 # Unknown component.
425 response = get_download_redirect(ref.id, "badcomponent")
426 self.assertEqual(response.status_code, 404)
428 # Not specifying the component for a multi-file "disassembled
429 # composite".
430 response = get_download_redirect(ref.id, None)
431 self.assertEqual(response.status_code, 422)
433 # Unknown dataset.
434 response = get_download_redirect(uuid.UUID("59467c1b-fa13-4f7a-8ff8-cd83e092e563"))
435 self.assertEqual(response.status_code, 404)
437 def test_auth_check(self):
438 # This is checking that the unit-test middleware for validating the
439 # authentication headers is working. It doesn't test actual server
440 # functionality -- in a real deployment, the authentication headers are
441 # handled by GafaelfawrIngress, not our app.
442 with self.assertRaises(UnhandledServerError):
443 self.client.get("/v1/dataset_type/int")
445 def test_exception_logging(self):
446 app = create_app()
448 def raise_error():
449 raise RuntimeError("An unhandled error")
451 app.dependency_overrides[butler_factory_dependency] = raise_error
452 client = TestClient(app, raise_server_exceptions=False)
454 with patch.object(safir.dependencies.logger, "logger_dependency") as mock_logger_dep:
455 mock_logger = NonCallableMock(["aerror"])
457 async def noop():
458 pass
460 mock_logger.aerror.return_value = noop()
462 async def get_logger():
463 return mock_logger
465 mock_logger_dep.return_value = get_logger()
466 client.get(
467 "/api/butler/repo/something/v1/dataset_type/int",
468 headers={"X-Auth-Request-User": "user-name", "X-Butler-Client-Request-Id": "request-id"},
469 )
470 mock_logger_dep.assert_called_once()
472 mock_logger.aerror.assert_called_once()
473 args, kwargs = mock_logger.aerror.call_args
474 self.assertIsInstance(kwargs["exc_info"], RuntimeError)
475 self.assertEqual(kwargs["clientRequestId"], "request-id")
476 self.assertEqual(kwargs["user"], "user-name")
478 def test_query_keepalive(self):
479 """Test that long-running queries stream keep-alive messages to stop
480 the HTTP connection from closing before they are able to return
481 results.
482 """
483 # Normally it takes 15 seconds for a timeout -- mock it to trigger
484 # immediately instead.
485 with patch.object(
486 lsst.daf.butler.remote_butler.server.handlers._query_streaming, "_timeout"
487 ) as mock_timeout:
488 # Hook into QueryDriver to track the number of keep-alives we have
489 # seen.
490 with patch.object(
491 lsst.daf.butler.remote_butler._query_results, "_received_keep_alive"
492 ) as mock_keep_alive:
493 mock_timeout.side_effect = _timeout_twice()
494 with self.butler.query() as query:
495 datasets = list(query.datasets("bias", "imported_g"))
496 self.assertEqual(len(datasets), 3)
497 self.assertGreaterEqual(mock_timeout.call_count, 3)
498 self.assertGreaterEqual(mock_keep_alive.call_count, 2)
500 def test_query_retries(self):
501 """Test that the server will send HTTP status 503 to put backpressure
502 on clients if it is overloaded, and that the client will retry if this
503 happens.
504 """
505 query_event = threading.Event()
506 retry_event = asyncio.Event()
508 async def block_first_request() -> None:
509 # Signal the unit tests that we have reached the critical section
510 # in the server, where the first client has reserved the query
511 # slot.
512 query_event.set()
513 # Block inside the query, until the 2nd client has been forced to
514 # retry.
515 await retry_event.wait()
517 async def block_second_request() -> None:
518 # Release the first client, so it can finish its query and prevent
519 # this client from being blocked on the next go-round.
520 retry_event.set()
522 def do_query(butler: Butler) -> list[DatasetRef]:
523 return butler.query_datasets("bias", "imported_g")
525 with (
526 patch.object(
527 lsst.daf.butler.remote_butler.server.handlers._query_limits,
528 "_MAXIMUM_CONCURRENT_STREAMING_QUERIES",
529 new=1,
530 ),
531 patch.object(
532 lsst.daf.butler.remote_butler.server.handlers._query_limits, "_QUERY_RETRY_SECONDS", new=1
533 ),
534 patch.object(
535 lsst.daf.butler.remote_butler.server.handlers._query_limits,
536 "_block_query_for_unit_test",
537 new=AsyncMock(wraps=block_first_request),
538 ) as mock_first_client,
539 patch.object(
540 lsst.daf.butler.remote_butler.server.handlers._query_limits,
541 "_block_retry_for_unit_test",
542 new=AsyncMock(wraps=block_second_request),
543 ) as mock_second_client,
544 ThreadPoolExecutor(max_workers=1) as exec1,
545 ThreadPoolExecutor(max_workers=1) as exec2,
546 ):
547 first_butler = self.butler
548 second_butler = self.butler.clone()
550 # Run the first client up until the server starts executing its
551 # query.
552 future1 = exec1.submit(do_query, first_butler)
553 event_reached = query_event.wait(60)
554 if not event_reached:
555 raise TimeoutError("Server did not execute query logic as expected.")
557 # Start the second client, which will trigger the retry logic and
558 # release the first client to finish its query.
559 future2 = exec2.submit(do_query, second_butler)
561 result1 = future1.result(60)
562 result2 = future2.result(60)
563 self.assertEqual(len(result1), 3)
564 self.assertEqual(len(result2), 3)
565 # The original thread should have gone through this section, and
566 # then the 2nd thread after it retries.
567 self.assertEqual(mock_first_client.await_count, 2)
568 # We should have triggered the retry logic at least once, but it
569 # might occur multiple times depending how long the first client
570 # takes to finish.
571 self.assertGreaterEqual(mock_second_client.await_count, 1)
573 # TODO DM-46204: This can be removed once the RSP recommended image has
574 # been upgraded to a version that contains DM-46129.
575 def test_deprecated_collection_endpoints(self):
576 # These REST endpoints are no longer used by Butler client so they need
577 # to be checked separately until they can be removed.
578 json = self.butler._connection.get(
579 "collection_info",
580 params={"name": "imported_g", "include_doc": True, "include_parents": True},
581 ).json()
582 self.assertEqual(json["name"], "imported_g")
583 self.assertEqual(json["type"], 1)
585 json = self.butler._connection.post(
586 "query_collections",
587 QueryCollectionsRequestModel(
588 search=["imported_*"], collection_types=[1], flatten_chains=False, include_chains=False
589 ),
590 ).json()
591 self.assertCountEqual(json["collections"], ["imported_g", "imported_r"])
593 def test_oversized_data_coordinate_upload(self):
594 with self.butler.query() as query:
595 ref = self.simple_dataset_ref
596 data_id = ref.dataId
597 data_coordinates = [DataCoordinate.standardize(data_id, visit=x) for x in range(100_001)]
598 with self.assertRaisesRegex(InvalidQueryError, "data coordinate rows"):
599 list(query.join_data_coordinates(data_coordinates).datasets(ref.datasetType, ref.run))
602@unittest.skipIf(create_test_server is None, f"Server dependencies not installed: {reason_text}")
603class ButlerClientServerAuthorizationTestCase(unittest.TestCase):
604 """Test authentication/authorization functionality."""
606 def test_group_authorization(self):
607 """Test that group membership repository authorization is checked when
608 repository is accessed.
609 """
610 with create_test_server(TESTDIR) as server_instance:
611 mock = MockGafaelfawrGroupAuthorizer()
612 server_instance.app.dependency_overrides[authorizer_dependency] = lambda: mock
613 server_instance.direct_butler.registry.registerDatasetType(
614 DatasetType("bias", [], "int", universe=server_instance.direct_butler.dimensions)
615 )
616 server_instance.direct_butler.collections.register("collection")
617 butler = server_instance.remote_butler
618 mock.set_response(False)
619 with self.assertRaises(ButlerServerError) as e:
620 butler.get_dataset_type("bias")
621 self.assertEqual(e.exception.status_code, 403)
622 with self.assertRaises(ButlerServerError) as e:
623 butler.query_datasets("bias", collections="*", find_first=False)
624 self.assertEqual(e.exception.status_code, 403)
626 mock.set_response(True)
627 self.assertEqual(butler.get_dataset_type("bias").name, "bias")
628 self.assertEqual(butler.query_datasets("bias", collections="collection", explain=False), [])
630 def test_cadc_auth(self) -> None:
631 """Test server running in CADC auth mode."""
632 with mock_config() as config:
633 config.authentication = "cadc"
634 config.gafaelfawr_url = "DISABLED"
635 with create_test_server(TESTDIR, server_config=config) as instance:
636 self.assertIsInstance(instance.remote_butler._connection.auth, CadcAuthenticationProvider)
638 # Set up a dataset backed by an HTTP URL.
639 # CADC uses a plain HTTP service, not S3, for hosting Butler
640 # artifacts.
641 dataset_type = DatasetType("test", [], "int", universe=instance.direct_butler.dimensions)
642 ref = DatasetRef(
643 datasetType=dataset_type,
644 dataId=DataCoordinate.makeEmpty(instance.direct_butler.dimensions),
645 run="ingest/run",
646 )
647 path = ResourcePath("https://fake-server.example/some-directory/file.json")
648 dataset = FileDataset(path, ref)
649 # ingest() insists on doing file existence checks, and we don't
650 # have an HTTP server to point it at.
651 with unittest.mock.patch(
652 "lsst.daf.butler.datastores.fileDatastore.FileDatastore._standardizeIngestPath"
653 ) as mock:
654 mock.return_value = path
655 instance.direct_butler.ingest(dataset, transfer="direct", record_validation_info=False)
657 # At the CADC, paths used for file download should NOT be a
658 # signed URL, and should have authentication headers attached.
659 def check_path(path_to_check: ResourcePath):
660 self.assertEqual(str(path_to_check), str(path))
661 assert isinstance(path_to_check, HttpResourcePath)
662 self.assertIsNotNone(path_to_check._extra_headers)
663 self.assertIsNotNone(path_to_check._extra_headers.get("Authorization"))
665 check_path(instance.remote_butler.getURI(ref))
666 transfer_map = instance.remote_butler._file_transfer_source.get_file_info_for_transfer(
667 [ref.id]
668 )
669 check_path(transfer_map[ref.id][0].location.pathInStore)
672def _create_corrupted_dataset(repo: MetricTestRepo) -> DatasetRef:
673 run = "corrupted-run"
674 ref = repo.addDataset({"instrument": "DummyCamComp", "visit": 423}, run=run)
675 uris = repo.butler.getURIs(ref)
676 oneOfTheComponents = list(uris.componentURIs.values())[0]
677 oneOfTheComponents.write("corrupted data")
678 return ref
681def _create_simple_dataset(butler: Butler) -> DatasetRef:
682 dataset_type = addDatasetType(butler, "test_int", {"instrument", "visit"}, "int")
683 ref = butler.put(123, dataset_type, dataId={"instrument": "DummyCamComp", "visit": 423}, run="ingest/run")
684 return ref
687def _timeout_twice():
688 """Return a mock side-effect function that raises a timeout error the first
689 two times it is called.
690 """
691 count = 0
693 def timeout(*args):
694 nonlocal count
695 count += 1
696 if count <= 2:
697 raise TimeoutError()
698 return DEFAULT
700 return timeout
703@unittest.skipIf(create_test_server is None, f"Server dependencies not installed: {reason_text}")
704class QueryLimitsTestCase(unittest.IsolatedAsyncioTestCase):
705 """Test details of the code that limits the maximum number of concurrent
706 queries in the server.
707 """
709 async def test_query_limits(self):
710 limits = lsst.daf.butler.remote_butler.server.handlers._query_limits.QueryLimits()
712 await limits.enforce_query_limits("user1") # under limit, doesn't raise
713 async with limits.track_query("user1"):
714 await limits.enforce_query_limits("user1") # under limit, doesn't raise
715 async with limits.track_query("user1"):
716 with self.assertRaises(fastapi.HTTPException) as exc:
717 await limits.enforce_query_limits("user1")
718 self.assertEqual(exc.exception.status_code, 429)
721if __name__ == "__main__":
722 unittest.main()