Coverage for python/lsst/daf/butler/tests/utils.py: 99%
122 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-22 09:23 +0000
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
30from unittest.mock import patch
32__all__ = ()
34import os
35import shutil
36import tempfile
37from collections.abc import Callable, Iterator, Sequence
38from contextlib import contextmanager
39from typing import TYPE_CHECKING, Any
41import astropy
42from astropy.table import Table as AstropyTable
44from lsst.resources import ResourcePathExpression
46from .. import Butler, ButlerConfig, Config, DatasetRef, StorageClassFactory, Timespan
47from .._collection_type import CollectionType
48from ..datastore import NullDatastore
49from ..dimensions import DimensionConfig
50from ..direct_butler import DirectButler
51from ..registry.sql_registry import RegistryConfig, SqlRegistry
52from ..tests import MetricsExample, addDatasetType
53from ._repo_template_cache import make_repo_for_test
55if TYPE_CHECKING:
56 import unittest
58 from lsst.daf.butler import DatasetType
60 class TestCaseMixin(unittest.TestCase):
61 """Base class for mixin test classes that use TestCase methods."""
63 pass
65else:
67 class TestCaseMixin:
68 """Do-nothing definition of mixin base class for regular execution."""
70 pass
73def makeTestTempDir(default_base: str) -> str:
74 """Create a temporary directory for test usage.
76 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
77 environment variable is set, falling back to ``default_base`` if it is
78 not.
80 Parameters
81 ----------
82 default_base : `str`
83 Default parent directory.
85 Returns
86 -------
87 dir : `str`
88 Name of the new temporary directory.
89 """
90 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
91 return tempfile.mkdtemp(dir=base)
94def removeTestTempDir(root: str | None) -> None:
95 """Attempt to remove a temporary test directory, but do not raise if
96 unable to.
98 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
99 to ``shutil.rmtree`` at close, making it safe to use on NFS.
101 Parameters
102 ----------
103 root : `str`, optional
104 Name of the directory to be removed. If `None`, nothing will be done.
105 """
106 if root is not None and os.path.exists(root): 106 ↛ exitline 106 didn't return from function 'removeTestTempDir' because the condition on line 106 was always true
107 shutil.rmtree(root, ignore_errors=True)
110@contextmanager
111def safeTestTempDir(default_base: str) -> Iterator[str]:
112 """Return a context manager that creates a temporary directory and then
113 attempts to remove it.
115 Parameters
116 ----------
117 default_base : `str`
118 Default parent directory, forwarded to `makeTestTempDir`.
120 Returns
121 -------
122 context : `contextlib.ContextManager`
123 A context manager that returns the new directory name on ``__enter__``
124 and removes the temporary directory (via `removeTestTempDir`) on
125 ``__exit__``.
126 """
127 root = makeTestTempDir(default_base)
128 try:
129 yield root
130 finally:
131 removeTestTempDir(root)
134def create_populated_sqlite_registry(
135 *args: ResourcePathExpression,
136 registry_config: RegistryConfig | None = None,
137 dimension_config: DimensionConfig | None = None,
138) -> Butler:
139 """Create an in-memory registry-only sqlite butler and populate it.
141 Parameters
142 ----------
143 *args : convertible to `lsst.resources.ResourcePath`
144 Paths to export YAML files that should be imported.
145 registry_config : ``RegistryConfig``, optional
146 Registry configuration to use as the basis for the Butler
147 configuration.
148 dimension_config : ``DimensionConfig``, optional
149 Dimension universe configuration.
151 Returns
152 -------
153 butler : `Butler`
154 New butler populated with the specified import files.
155 """
156 config = ButlerConfig()
157 if registry_config is not None:
158 config["registry"] = registry_config
159 config[".registry.db"] = "sqlite://"
160 registry = SqlRegistry.createFromConfig(config["registry"], dimension_config)
161 butler = DirectButler(
162 config=config,
163 registry=registry,
164 datastore=NullDatastore(None, None),
165 storageClasses=StorageClassFactory(),
166 )
167 for arg in args:
168 butler.import_(filename=arg, without_datastore=True)
169 return butler
172class ButlerTestHelper:
173 """Mixin with helpers for unit tests."""
175 assertEqual: Callable
176 assertIsInstance: Callable
177 maxDiff: int | None
179 def assertAstropyTablesEqual(
180 self,
181 tables: AstropyTable | Sequence[AstropyTable],
182 expectedTables: AstropyTable | Sequence[AstropyTable],
183 filterColumns: bool = False,
184 unorderedRows: bool = False,
185 ) -> None:
186 """Verify that a list of astropy tables matches a list of expected
187 astropy tables.
189 Parameters
190 ----------
191 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
192 The table or tables that should match the expected tables.
193 expectedTables : `astropy.table.Table`
194 or iterable [`astropy.table.Table`]
195 The tables with expected values to which the tables under test will
196 be compared.
197 filterColumns : `bool`
198 If `True` then only compare columns that exist in
199 ``expectedTables``.
200 unorderedRows : `bool`, optional
201 If `True` (`False` is default), don't require tables to have their
202 rows in the same order.
203 """
204 # If a single table is passed in for tables or expectedTables, put it
205 # in a list.
206 if isinstance(tables, AstropyTable):
207 tables = [tables]
208 if isinstance(expectedTables, AstropyTable):
209 expectedTables = [expectedTables]
210 self.assertEqual(len(tables), len(expectedTables))
211 for table, expected in zip(tables, expectedTables, strict=True):
212 # Assert that we are testing what we think we are testing:
213 self.assertIsInstance(table, AstropyTable)
214 self.assertIsInstance(expected, AstropyTable)
215 if filterColumns:
216 table = table.copy()
217 table.keep_columns(expected.colnames)
218 if unorderedRows:
219 table = table.copy()
220 table.sort(table.colnames)
221 expected = expected.copy()
222 expected.sort(expected.colnames)
223 # Assert that they match.
224 # Recommendation from Astropy Slack is to format the table into
225 # lines for comparison. We do not compare column data types.
226 table1 = table.pformat()
227 expected1 = expected.pformat()
228 original_max = self.maxDiff
229 self.maxDiff = None # This is required to get the full diff.
230 try:
231 self.assertEqual(table1, expected1, f"Table:\n{table}\n\nvs Expected:\n{expected}")
232 finally:
233 self.maxDiff = original_max
236def readTable(textTable: str) -> AstropyTable:
237 """Read an astropy table from formatted text.
239 Contains formatting that causes the astropy table to print an empty string
240 instead of "--" for missing/unpopulated values in the text table.
242 Parameters
243 ----------
244 textTable : `str`
245 The text version of the table to read.
247 Returns
248 -------
249 table : `astropy.table.Table`
250 The table as an astropy table.
251 """
252 return AstropyTable.read(
253 textTable,
254 format="ascii",
255 data_start=2, # skip the header row and the header row underlines.
256 fill_values=[("", 0, "")],
257 )
260class MetricTestRepo:
261 """Creates and manage a test repository on disk with datasets that
262 may be queried and modified for unit tests.
264 Parameters
265 ----------
266 root : `str`
267 The location of the repository, to pass to ``Butler.makeRepo``.
268 configFile : `str`
269 The path to the config file, to pass to ``Butler.makeRepo``.
270 forceConfigRoot : `bool`, optional
271 If `False`, any values present in the supplied ``config`` that
272 would normally be reset are not overridden and will appear
273 directly in the output config. Passed to ``Butler.makeRepo``.
274 storageClassName : `bool` or `None`, optional
275 Name of storage class to use for datasets added to the test repository.
276 A default will be used if none is specified.
277 """
279 METRICS_EXAMPLE_SUMMARY = {"AM1": 5.2, "AM2": 30.6}
280 """The summary data included in ``MetricsExample`` objects stored in the
281 test repo
282 """
284 _DEFAULT_RUN = "ingest/run"
285 _DEFAULT_TAG = "ingest"
286 _DEFAULT_STORAGE_CLASS = "StructuredCompositeReadComp"
288 @staticmethod
289 def _makeExampleMetrics() -> MetricsExample:
290 """Make an object to put into the repository."""
291 return MetricsExample(
292 MetricTestRepo.METRICS_EXAMPLE_SUMMARY,
293 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
294 [563, 234, 456.7, 752, 8, 9, 27],
295 )
297 def __init__(
298 self,
299 root: str,
300 configFile: str,
301 forceConfigRoot: bool = True,
302 storageClassName: str | None = None,
303 ) -> None:
304 self.root = root
305 butlerConfigFile = make_repo_for_test(
306 self.root, config=Config(configFile), forceConfigRoot=forceConfigRoot
307 )
308 butler = Butler.from_config(butlerConfigFile, run=self._DEFAULT_RUN, collections=[self._DEFAULT_TAG])
309 self._do_init(butler, butlerConfigFile, storageClassName)
311 @classmethod
312 def create_from_butler(
313 cls, butler: Butler, butler_config_file: str | Config, storageClassName: str | None = None
314 ) -> MetricTestRepo:
315 """Create a MetricTestRepo from an existing Butler instance.
317 Parameters
318 ----------
319 butler
320 `Butler` instance used for setting up the repository.
321 butler_config_file
322 Path to the config file or the `Config` instance used to set up
323 that Butler instance.
324 storageClassName
325 Name of storage class to use for datasets added to the test
326 repository. A default will be used if none is specified.
328 Returns
329 -------
330 repo
331 New instance of `MetricTestRepo` using the provided `Butler`
332 instance.
333 """
334 self = cls.__new__(cls)
335 butler = butler.clone(run=self._DEFAULT_RUN, collections=[self._DEFAULT_TAG])
336 self._do_init(butler, butler_config_file, storageClassName)
337 return self
339 def _do_init(
340 self, butler: Butler, butlerConfigFile: str | Config, storageClassName: str | None = None
341 ) -> None:
342 self.butler = butler
343 self.storageClassFactory = StorageClassFactory()
344 self.storageClassFactory.addFromConfig(butlerConfigFile)
346 # New datasets will be added to run and tag, but we will only look in
347 # tag when looking up datasets.
348 self.butler.collections.register(self._DEFAULT_TAG, CollectionType.TAGGED)
350 if storageClassName is None:
351 storageClassName = self._DEFAULT_STORAGE_CLASS
353 # Create and register a DatasetType
354 self.datasetType = addDatasetType(
355 self.butler, "test_metric_comp", {"instrument", "visit"}, storageClassName
356 )
358 # Add needed Dimensions
359 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
360 self.butler.registry.insertDimensionData(
361 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
362 )
363 self.butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101})
364 self.butler.registry.insertDimensionData(
365 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
366 )
367 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
368 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
369 self.butler.registry.insertDimensionData(
370 "visit",
371 dict(
372 instrument="DummyCamComp",
373 id=423,
374 name="fourtwentythree",
375 physical_filter="d-r",
376 timespan=Timespan(visitStart, visitEnd),
377 day_obs=20200101,
378 ),
379 )
380 self.butler.registry.insertDimensionData(
381 "visit",
382 dict(
383 instrument="DummyCamComp",
384 id=424,
385 name="fourtwentyfour",
386 physical_filter="d-r",
387 day_obs=20200101,
388 ),
389 )
391 self.ref1 = self.addDataset({"instrument": "DummyCamComp", "visit": 423})
392 self.ref2 = self.addDataset({"instrument": "DummyCamComp", "visit": 424})
394 def addDataset(
395 self, dataId: dict[str, Any], run: str | None = None, datasetType: DatasetType | None = None
396 ) -> DatasetRef:
397 """Create a new example metric and add it to the named run with the
398 given dataId.
400 Overwrites tags, so this does not try to associate the new dataset with
401 existing tags. (If/when tags are needed this can be added to the
402 arguments of this function.)
404 Parameters
405 ----------
406 dataId : `dict`
407 The dataId for the new metric.
408 run : `str`, optional
409 The name of the run to create and add a dataset to. If `None`, the
410 dataset will be added to the root butler.
411 datasetType : ``DatasetType``, optional
412 The dataset type of the added dataset. If `None`, will use the
413 default dataset type.
415 Returns
416 -------
417 datasetRef : `DatasetRef`
418 A reference to the added dataset.
419 """
420 if run:
421 self.butler.collections.register(run)
422 else:
423 run = self._DEFAULT_RUN
424 metric = self._makeExampleMetrics()
425 return self.butler.put(
426 metric, self.datasetType if datasetType is None else datasetType, dataId, run=run
427 )
430@contextmanager
431def mock_env(new_environment: dict[str, str]) -> Iterator[None]:
432 """Context manager to clear the process environment variables, replace them
433 with new values, and restore them at the end of the test.
435 Parameters
436 ----------
437 new_environment : `dict` [`str`, `str`]
438 New environment variable values.
439 """
440 with patch.dict(os.environ, new_environment, clear=True):
441 yield