Coverage for python/lsst/daf/butler/tests/utils.py: 36%
93 statements
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-01 11:20 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-01 11:20 +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 .. import Butler, Config, DatasetRef, StorageClassFactory, Timespan
45from ..registry import CollectionType
46from ..tests import MetricsExample, addDatasetType
48if TYPE_CHECKING:
49 import unittest
51 from lsst.daf.butler import DatasetType
53 class TestCaseMixin(unittest.TestCase):
54 """Base class for mixin test classes that use TestCase methods."""
56 pass
58else:
60 class TestCaseMixin:
61 """Do-nothing definition of mixin base class for regular execution."""
63 pass
66def makeTestTempDir(default_base: str) -> str:
67 """Create a temporary directory for test usage.
69 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
70 environment variable is set, falling back to ``default_base`` if it is
71 not.
73 Parameters
74 ----------
75 default_base : `str`
76 Default parent directory.
78 Returns
79 -------
80 dir : `str`
81 Name of the new temporary directory.
82 """
83 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
84 return tempfile.mkdtemp(dir=base)
87def removeTestTempDir(root: str | None) -> None:
88 """Attempt to remove a temporary test directory, but do not raise if
89 unable to.
91 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
92 to ``shutil.rmtree`` at close, making it safe to use on NFS.
94 Parameters
95 ----------
96 root : `str`, optional
97 Name of the directory to be removed. If `None`, nothing will be done.
98 """
99 if root is not None and os.path.exists(root):
100 shutil.rmtree(root, ignore_errors=True)
103@contextmanager
104def safeTestTempDir(default_base: str) -> Iterator[str]:
105 """Return a context manager that creates a temporary directory and then
106 attempts to remove it.
108 Parameters
109 ----------
110 default_base : `str`
111 Default parent directory, forwarded to `makeTestTempDir`.
113 Returns
114 -------
115 context : `contextlib.ContextManager`
116 A context manager that returns the new directory name on ``__enter__``
117 and removes the temporary directory (via `removeTestTempDir`) on
118 ``__exit__``.
119 """
120 root = makeTestTempDir(default_base)
121 try:
122 yield root
123 finally:
124 removeTestTempDir(root)
127class ButlerTestHelper:
128 """Mixin with helpers for unit tests."""
130 assertEqual: Callable
131 assertIsInstance: Callable
132 maxDiff: int | None
134 def assertAstropyTablesEqual(
135 self,
136 tables: AstropyTable | Sequence[AstropyTable],
137 expectedTables: AstropyTable | Sequence[AstropyTable],
138 filterColumns: bool = False,
139 unorderedRows: bool = False,
140 ) -> None:
141 """Verify that a list of astropy tables matches a list of expected
142 astropy tables.
144 Parameters
145 ----------
146 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
147 The table or tables that should match the expected tables.
148 expectedTables : `astropy.table.Table`
149 or iterable [`astropy.table.Table`]
150 The tables with expected values to which the tables under test will
151 be compared.
152 filterColumns : `bool`
153 If `True` then only compare columns that exist in
154 ``expectedTables``.
155 unorderedRows : `bool`, optional
156 If `True` (`False` is default), don't require tables to have their
157 rows in the same order.
158 """
159 # If a single table is passed in for tables or expectedTables, put it
160 # in a list.
161 if isinstance(tables, AstropyTable):
162 tables = [tables]
163 if isinstance(expectedTables, AstropyTable):
164 expectedTables = [expectedTables]
165 self.assertEqual(len(tables), len(expectedTables))
166 for table, expected in zip(tables, expectedTables, strict=True):
167 # Assert that we are testing what we think we are testing:
168 self.assertIsInstance(table, AstropyTable)
169 self.assertIsInstance(expected, AstropyTable)
170 if filterColumns:
171 table = table.copy()
172 table.keep_columns(expected.colnames)
173 if unorderedRows:
174 table = table.copy()
175 table.sort(table.colnames)
176 expected = expected.copy()
177 expected.sort(expected.colnames)
178 # Assert that they match.
179 # Recommendation from Astropy Slack is to format the table into
180 # lines for comparison. We do not compare column data types.
181 table1 = table.pformat_all()
182 expected1 = expected.pformat_all()
183 original_max = self.maxDiff
184 self.maxDiff = None # This is required to get the full diff.
185 try:
186 self.assertEqual(table1, expected1)
187 finally:
188 self.maxDiff = original_max
191def readTable(textTable: str) -> AstropyTable:
192 """Read an astropy table from formatted text.
194 Contains formatting that causes the astropy table to print an empty string
195 instead of "--" for missing/unpopulated values in the text table.
197 Parameters
198 ----------
199 textTable : `str`
200 The text version of the table to read.
202 Returns
203 -------
204 table : `astropy.table.Table`
205 The table as an astropy table.
206 """
207 return AstropyTable.read(
208 textTable,
209 format="ascii",
210 data_start=2, # skip the header row and the header row underlines.
211 fill_values=[("", 0, "")],
212 )
215class MetricTestRepo:
216 """Creates and manage a test repository on disk with datasets that
217 may be queried and modified for unit tests.
219 Parameters
220 ----------
221 root : `str`
222 The location of the repository, to pass to ``Butler.makeRepo``.
223 configFile : `str`
224 The path to the config file, to pass to ``Butler.makeRepo``.
225 forceConfigRoot : `bool`, optional
226 If `False`, any values present in the supplied ``config`` that
227 would normally be reset are not overridden and will appear
228 directly in the output config. Passed to ``Butler.makeRepo``.
229 """
231 METRICS_EXAMPLE_SUMMARY = {"AM1": 5.2, "AM2": 30.6}
232 """The summary data included in ``MetricsExample`` objects stored in the
233 test repo
234 """
236 @staticmethod
237 def _makeExampleMetrics() -> MetricsExample:
238 """Make an object to put into the repository."""
239 return MetricsExample(
240 MetricTestRepo.METRICS_EXAMPLE_SUMMARY,
241 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
242 [563, 234, 456.7, 752, 8, 9, 27],
243 )
245 def __init__(self, root: str, configFile: str, forceConfigRoot: bool = True) -> None:
246 self.root = root
247 Butler.makeRepo(self.root, config=Config(configFile), forceConfigRoot=forceConfigRoot)
248 butlerConfigFile = os.path.join(self.root, "butler.yaml")
249 self.storageClassFactory = StorageClassFactory()
250 self.storageClassFactory.addFromConfig(butlerConfigFile)
252 # New datasets will be added to run and tag, but we will only look in
253 # tag when looking up datasets.
254 run = "ingest/run"
255 tag = "ingest"
256 self.butler = Butler.from_config(butlerConfigFile, run=run, collections=[tag])
257 self.butler.registry.registerCollection(tag, CollectionType.TAGGED)
259 # Create and register a DatasetType
260 self.datasetType = addDatasetType(
261 self.butler, "test_metric_comp", {"instrument", "visit"}, "StructuredCompositeReadComp"
262 )
264 # Add needed Dimensions
265 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
266 self.butler.registry.insertDimensionData(
267 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
268 )
269 self.butler.registry.insertDimensionData(
270 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
271 )
272 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
273 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
274 self.butler.registry.insertDimensionData(
275 "visit",
276 dict(
277 instrument="DummyCamComp",
278 id=423,
279 name="fourtwentythree",
280 physical_filter="d-r",
281 timespan=Timespan(visitStart, visitEnd),
282 ),
283 )
284 self.butler.registry.insertDimensionData(
285 "visit",
286 dict(
287 instrument="DummyCamComp",
288 id=424,
289 name="fourtwentyfour",
290 physical_filter="d-r",
291 ),
292 )
294 self.addDataset({"instrument": "DummyCamComp", "visit": 423})
295 self.addDataset({"instrument": "DummyCamComp", "visit": 424})
297 def addDataset(
298 self, dataId: dict[str, Any], run: str | None = None, datasetType: DatasetType | None = None
299 ) -> DatasetRef:
300 """Create a new example metric and add it to the named run with the
301 given dataId.
303 Overwrites tags, so this does not try to associate the new dataset with
304 existing tags. (If/when tags are needed this can be added to the
305 arguments of this function.)
307 Parameters
308 ----------
309 dataId : `dict`
310 The dataId for the new metric.
311 run : `str`, optional
312 The name of the run to create and add a dataset to. If `None`, the
313 dataset will be added to the root butler.
314 datasetType : ``DatasetType``, optional
315 The dataset type of the added dataset. If `None`, will use the
316 default dataset type.
318 Returns
319 -------
320 datasetRef : `DatasetRef`
321 A reference to the added dataset.
322 """
323 if run:
324 self.butler.registry.registerCollection(run, type=CollectionType.RUN)
325 metric = self._makeExampleMetrics()
326 return self.butler.put(
327 metric, self.datasetType if datasetType is None else datasetType, dataId, run=run
328 )
331@contextmanager
332def mock_env(new_environment: dict[str, str]) -> Iterator[None]:
333 """Context manager to clear the process environment variables, replace them
334 with new values, and restore them at the end of the test.
336 Parameters
337 ----------
338 new_environment : `dict`[`str`, `str`]
339 New environment variable values.
340 """
341 with patch.dict(os.environ, new_environment, clear=True):
342 yield