Coverage for python/lsst/daf/butler/tests/utils.py: 25%
87 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-17 02:01 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-17 02:01 -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 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 contextlib import contextmanager
30from typing import Optional
32import astropy
33from astropy.table import Table as AstropyTable
35from .. import Butler, Config, StorageClassFactory
36from ..registry import CollectionType
37from ..tests import MetricsExample, addDatasetType
40def makeTestTempDir(default_base: str) -> str:
41 """Create a temporary directory for test usage.
43 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
44 environment variable is set, falling back to ``default_base`` if it is
45 not.
47 Parameters
48 ----------
49 default_base : `str`
50 Default parent directory.
52 Returns
53 -------
54 dir : `str`
55 Name of the new temporary directory.
56 """
57 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
58 return tempfile.mkdtemp(dir=base)
61def removeTestTempDir(root: Optional[str]) -> None:
62 """Attempt to remove a temporary test directory, but do not raise if
63 unable to.
65 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
66 to ``shutil.rmtree`` at close, making it safe to use on NFS.
68 Parameters
69 ----------
70 root : `str`, optional
71 Name of the directory to be removed. If `None`, nothing will be done.
72 """
73 if root is not None and os.path.exists(root):
74 shutil.rmtree(root, ignore_errors=True)
77@contextmanager
78def safeTestTempDir(default_base: str) -> str:
79 """Return a context manager that creates a temporary directory and then
80 attempts to remove it.
82 Parameters
83 ----------
84 default_base : `str`
85 Default parent directory, forwarded to `makeTestTempDir`.
87 Returns
88 -------
89 context : `contextlib.ContextManager`
90 A context manager that returns the new directory name on ``__enter__``
91 and removes the temporary directory (via `removeTestTempDir`) on
92 ``__exit__``.
93 """
94 root = makeTestTempDir(default_base)
95 try:
96 yield root
97 finally:
98 removeTestTempDir(root)
101class ButlerTestHelper:
102 """Mixin with helpers for unit tests."""
104 def assertAstropyTablesEqual(self, tables, expectedTables, filterColumns=False, unorderedRows=False):
105 """Verify that a list of astropy tables matches a list of expected
106 astropy tables.
108 Parameters
109 ----------
110 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
111 The table or tables that should match the expected tables.
112 expectedTables : `astropy.table.Table`
113 or iterable [`astropy.table.Table`]
114 The tables with expected values to which the tables under test will
115 be compared.
116 filterColumns : `bool`
117 If `True` then only compare columns that exist in
118 ``expectedTables``.
119 unorderedRows : `bool`, optional
120 If `True` (`False` is default), don't require tables to have their
121 rows in the same order.
122 """
123 # If a single table is passed in for tables or expectedTables, put it
124 # in a list.
125 if isinstance(tables, AstropyTable):
126 tables = [tables]
127 if isinstance(expectedTables, AstropyTable):
128 expectedTables = [expectedTables]
129 self.assertEqual(len(tables), len(expectedTables))
130 for table, expected in zip(tables, expectedTables):
131 # Assert that we are testing what we think we are testing:
132 self.assertIsInstance(table, AstropyTable)
133 self.assertIsInstance(expected, AstropyTable)
134 if filterColumns:
135 table = table.copy()
136 table.keep_columns(expected.colnames)
137 if unorderedRows:
138 table = table.copy()
139 table.sort(table.colnames)
140 expected = expected.copy()
141 expected.sort(expected.colnames)
142 # Assert that they match.
143 # Recommendation from Astropy Slack is to format the table into
144 # lines for comparison. We do not compare column data types.
145 table1 = table.pformat_all()
146 expected1 = expected.pformat_all()
147 original_max = self.maxDiff
148 self.maxDiff = None # This is required to get the full diff.
149 try:
150 self.assertEqual(table1, expected1)
151 finally:
152 self.maxDiff = original_max
155def readTable(textTable):
156 """Read an astropy table from formatted text.
158 Contains formatting that causes the astropy table to print an empty string
159 instead of "--" for missing/unpopulated values in the text table.
162 Parameters
163 ----------
164 textTable : `str`
165 The text version of the table to read.
167 Returns
168 -------
169 table : `astropy.table.Table`
170 The table as an astropy table.
171 """
172 return AstropyTable.read(
173 textTable,
174 format="ascii",
175 data_start=2, # skip the header row and the header row underlines.
176 fill_values=[("", 0, "")],
177 )
180class MetricTestRepo:
181 """Creates and manage a test repository on disk with datasets that
182 may be queried and modified for unit tests.
184 Parameters
185 ----------
186 root : `str`
187 The location of the repository, to pass to ``Butler.makeRepo``.
188 configFile : `str`
189 The path to the config file, to pass to ``Butler.makeRepo``.
190 """
192 @staticmethod
193 def _makeExampleMetrics():
194 """Make an object to put into the repository."""
195 return MetricsExample(
196 {"AM1": 5.2, "AM2": 30.6},
197 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
198 [563, 234, 456.7, 752, 8, 9, 27],
199 )
201 @staticmethod
202 def _makeDimensionData(id, name, datetimeBegin=None, datetimeEnd=None):
203 """Make a dict of dimensional data with default values to insert into
204 the registry.
205 """
206 data = dict(instrument="DummyCamComp", id=id, name=name, physical_filter="d-r", visit_system=1)
207 if datetimeBegin:
208 data["datetime_begin"] = datetimeBegin
209 data["datetime_end"] = datetimeEnd
210 return data
212 def __init__(self, root, configFile):
213 self.root = root
214 Butler.makeRepo(self.root, config=Config(configFile))
215 butlerConfigFile = os.path.join(self.root, "butler.yaml")
216 self.storageClassFactory = StorageClassFactory()
217 self.storageClassFactory.addFromConfig(butlerConfigFile)
219 # New datasets will be added to run and tag, but we will only look in
220 # tag when looking up datasets.
221 run = "ingest/run"
222 tag = "ingest"
223 self.butler = Butler(butlerConfigFile, run=run, collections=[tag])
224 self.butler.registry.registerCollection(tag, CollectionType.TAGGED)
226 # Create and register a DatasetType
227 self.datasetType = addDatasetType(
228 self.butler, "test_metric_comp", ("instrument", "visit"), "StructuredCompositeReadComp"
229 )
231 # Add needed Dimensions
232 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
233 self.butler.registry.insertDimensionData(
234 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
235 )
236 self.butler.registry.insertDimensionData(
237 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
238 )
239 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
240 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
241 self.butler.registry.insertDimensionData(
242 "visit",
243 dict(
244 instrument="DummyCamComp",
245 id=423,
246 name="fourtwentythree",
247 physical_filter="d-r",
248 datetimeBegin=visitStart,
249 datetimeEnd=visitEnd,
250 ),
251 )
252 self.butler.registry.insertDimensionData(
253 "visit",
254 dict(
255 instrument="DummyCamComp",
256 id=424,
257 name="fourtwentyfour",
258 physical_filter="d-r",
259 ),
260 )
262 self.addDataset({"instrument": "DummyCamComp", "visit": 423})
263 self.addDataset({"instrument": "DummyCamComp", "visit": 424})
265 def addDataset(self, dataId, run=None, datasetType=None):
266 """Create a new example metric and add it to the named run with the
267 given dataId.
269 Overwrites tags, so this does not try to associate the new dataset with
270 existing tags. (If/when tags are needed this can be added to the
271 arguments of this function.)
273 Parameters
274 ----------
275 dataId : `dict`
276 The dataId for the new metric.
277 run : `str`, optional
278 The name of the run to create and add a dataset to. If `None`, the
279 dataset will be added to the root butler.
280 datasetType : ``DatasetType``, optional
281 The dataset type of the added dataset. If `None`, will use the
282 default dataset type.
283 """
284 if run:
285 self.butler.registry.registerCollection(run, type=CollectionType.RUN)
286 metric = self._makeExampleMetrics()
287 self.butler.put(metric, self.datasetType if datasetType is None else datasetType, dataId, run=run)