Coverage for python/lsst/daf/butler/tests/utils.py: 30%
86 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-15 09:13 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-15 09:13 +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 program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22from __future__ import annotations
24__all__ = ()
26import os
27import shutil
28import tempfile
29from collections.abc import Callable, Iterator, Sequence
30from contextlib import contextmanager
31from typing import TYPE_CHECKING, Any
33import astropy
34from astropy.table import Table as AstropyTable
36from .. import Butler, Config, StorageClassFactory
37from ..registry import CollectionType
38from ..tests import MetricsExample, addDatasetType
40if TYPE_CHECKING:
41 import unittest
43 from lsst.daf.butler import DatasetType
45 class TestCaseMixin(unittest.TestCase):
46 """Base class for mixin test classes that use TestCase methods."""
48 pass
50else:
52 class TestCaseMixin:
53 """Do-nothing definition of mixin base class for regular execution."""
55 pass
58def makeTestTempDir(default_base: str) -> str:
59 """Create a temporary directory for test usage.
61 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
62 environment variable is set, falling back to ``default_base`` if it is
63 not.
65 Parameters
66 ----------
67 default_base : `str`
68 Default parent directory.
70 Returns
71 -------
72 dir : `str`
73 Name of the new temporary directory.
74 """
75 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
76 return tempfile.mkdtemp(dir=base)
79def removeTestTempDir(root: str | None) -> None:
80 """Attempt to remove a temporary test directory, but do not raise if
81 unable to.
83 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
84 to ``shutil.rmtree`` at close, making it safe to use on NFS.
86 Parameters
87 ----------
88 root : `str`, optional
89 Name of the directory to be removed. If `None`, nothing will be done.
90 """
91 if root is not None and os.path.exists(root):
92 shutil.rmtree(root, ignore_errors=True)
95@contextmanager
96def safeTestTempDir(default_base: str) -> Iterator[str]:
97 """Return a context manager that creates a temporary directory and then
98 attempts to remove it.
100 Parameters
101 ----------
102 default_base : `str`
103 Default parent directory, forwarded to `makeTestTempDir`.
105 Returns
106 -------
107 context : `contextlib.ContextManager`
108 A context manager that returns the new directory name on ``__enter__``
109 and removes the temporary directory (via `removeTestTempDir`) on
110 ``__exit__``.
111 """
112 root = makeTestTempDir(default_base)
113 try:
114 yield root
115 finally:
116 removeTestTempDir(root)
119class ButlerTestHelper:
120 """Mixin with helpers for unit tests."""
122 assertEqual: Callable
123 assertIsInstance: Callable
124 maxDiff: int | None
126 def assertAstropyTablesEqual(
127 self,
128 tables: AstropyTable | Sequence[AstropyTable],
129 expectedTables: AstropyTable | Sequence[AstropyTable],
130 filterColumns: bool = False,
131 unorderedRows: bool = False,
132 ) -> None:
133 """Verify that a list of astropy tables matches a list of expected
134 astropy tables.
136 Parameters
137 ----------
138 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
139 The table or tables that should match the expected tables.
140 expectedTables : `astropy.table.Table`
141 or iterable [`astropy.table.Table`]
142 The tables with expected values to which the tables under test will
143 be compared.
144 filterColumns : `bool`
145 If `True` then only compare columns that exist in
146 ``expectedTables``.
147 unorderedRows : `bool`, optional
148 If `True` (`False` is default), don't require tables to have their
149 rows in the same order.
150 """
151 # If a single table is passed in for tables or expectedTables, put it
152 # in a list.
153 if isinstance(tables, AstropyTable):
154 tables = [tables]
155 if isinstance(expectedTables, AstropyTable):
156 expectedTables = [expectedTables]
157 self.assertEqual(len(tables), len(expectedTables))
158 for table, expected in zip(tables, expectedTables):
159 # Assert that we are testing what we think we are testing:
160 self.assertIsInstance(table, AstropyTable)
161 self.assertIsInstance(expected, AstropyTable)
162 if filterColumns:
163 table = table.copy()
164 table.keep_columns(expected.colnames)
165 if unorderedRows:
166 table = table.copy()
167 table.sort(table.colnames)
168 expected = expected.copy()
169 expected.sort(expected.colnames)
170 # Assert that they match.
171 # Recommendation from Astropy Slack is to format the table into
172 # lines for comparison. We do not compare column data types.
173 table1 = table.pformat_all()
174 expected1 = expected.pformat_all()
175 original_max = self.maxDiff
176 self.maxDiff = None # This is required to get the full diff.
177 try:
178 self.assertEqual(table1, expected1)
179 finally:
180 self.maxDiff = original_max
183def readTable(textTable: str) -> AstropyTable:
184 """Read an astropy table from formatted text.
186 Contains formatting that causes the astropy table to print an empty string
187 instead of "--" for missing/unpopulated values in the text table.
190 Parameters
191 ----------
192 textTable : `str`
193 The text version of the table to read.
195 Returns
196 -------
197 table : `astropy.table.Table`
198 The table as an astropy table.
199 """
200 return AstropyTable.read(
201 textTable,
202 format="ascii",
203 data_start=2, # skip the header row and the header row underlines.
204 fill_values=[("", 0, "")],
205 )
208class MetricTestRepo:
209 """Creates and manage a test repository on disk with datasets that
210 may be queried and modified for unit tests.
212 Parameters
213 ----------
214 root : `str`
215 The location of the repository, to pass to ``Butler.makeRepo``.
216 configFile : `str`
217 The path to the config file, to pass to ``Butler.makeRepo``.
218 """
220 @staticmethod
221 def _makeExampleMetrics() -> MetricsExample:
222 """Make an object to put into the repository."""
223 return MetricsExample(
224 {"AM1": 5.2, "AM2": 30.6},
225 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
226 [563, 234, 456.7, 752, 8, 9, 27],
227 )
229 def __init__(self, root: str, configFile: str) -> None:
230 self.root = root
231 Butler.makeRepo(self.root, config=Config(configFile))
232 butlerConfigFile = os.path.join(self.root, "butler.yaml")
233 self.storageClassFactory = StorageClassFactory()
234 self.storageClassFactory.addFromConfig(butlerConfigFile)
236 # New datasets will be added to run and tag, but we will only look in
237 # tag when looking up datasets.
238 run = "ingest/run"
239 tag = "ingest"
240 self.butler = Butler(butlerConfigFile, run=run, collections=[tag])
241 self.butler.registry.registerCollection(tag, CollectionType.TAGGED)
243 # Create and register a DatasetType
244 self.datasetType = addDatasetType(
245 self.butler, "test_metric_comp", {"instrument", "visit"}, "StructuredCompositeReadComp"
246 )
248 # Add needed Dimensions
249 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
250 self.butler.registry.insertDimensionData(
251 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
252 )
253 self.butler.registry.insertDimensionData(
254 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
255 )
256 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
257 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
258 self.butler.registry.insertDimensionData(
259 "visit",
260 dict(
261 instrument="DummyCamComp",
262 id=423,
263 name="fourtwentythree",
264 physical_filter="d-r",
265 datetimeBegin=visitStart,
266 datetimeEnd=visitEnd,
267 ),
268 )
269 self.butler.registry.insertDimensionData(
270 "visit",
271 dict(
272 instrument="DummyCamComp",
273 id=424,
274 name="fourtwentyfour",
275 physical_filter="d-r",
276 ),
277 )
279 self.addDataset({"instrument": "DummyCamComp", "visit": 423})
280 self.addDataset({"instrument": "DummyCamComp", "visit": 424})
282 def addDataset(
283 self, dataId: dict[str, Any], run: str | None = None, datasetType: DatasetType | None = None
284 ) -> None:
285 """Create a new example metric and add it to the named run with the
286 given dataId.
288 Overwrites tags, so this does not try to associate the new dataset with
289 existing tags. (If/when tags are needed this can be added to the
290 arguments of this function.)
292 Parameters
293 ----------
294 dataId : `dict`
295 The dataId for the new metric.
296 run : `str`, optional
297 The name of the run to create and add a dataset to. If `None`, the
298 dataset will be added to the root butler.
299 datasetType : ``DatasetType``, optional
300 The dataset type of the added dataset. If `None`, will use the
301 default dataset type.
302 """
303 if run:
304 self.butler.registry.registerCollection(run, type=CollectionType.RUN)
305 metric = self._makeExampleMetrics()
306 self.butler.put(metric, self.datasetType if datasetType is None else datasetType, dataId, run=run)