Coverage for python/lsst/daf/butler/tests/utils.py: 33%
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
Shortcuts on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
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 io
27import os
28import shutil
29import tempfile
30from contextlib import contextmanager
31from typing import Optional
33import astropy
34from astropy.table import Table as AstropyTable
35from astropy.utils.diff import report_diff_values
37from .. import Butler, Config, StorageClassFactory
38from ..registry import CollectionType
39from ..tests import MetricsExample, addDatasetType
42def makeTestTempDir(default_base: str) -> str:
43 """Create a temporary directory for test usage.
45 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
46 environment variable is set, falling back to ``default_base`` if it is
47 not.
49 Parameters
50 ----------
51 default_base : `str`
52 Default parent directory.
54 Returns
55 -------
56 dir : `str`
57 Name of the new temporary directory.
58 """
59 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
60 return tempfile.mkdtemp(dir=base)
63def removeTestTempDir(root: Optional[str]) -> None:
64 """Attempt to remove a temporary test directory, but do not raise if
65 unable to.
67 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
68 to ``shutil.rmtree`` at close, making it safe to use on NFS.
70 Parameters
71 ----------
72 root : `str`, optional
73 Name of the directory to be removed. If `None`, nothing will be done.
74 """
75 if root is not None and os.path.exists(root):
76 shutil.rmtree(root, ignore_errors=True)
79@contextmanager
80def safeTestTempDir(default_base: str) -> str:
81 """Return a context manager that creates a temporary directory and then
82 attempts to remove it.
84 Parameters
85 ----------
86 default_base : `str`
87 Default parent directory, forwarded to `makeTestTempDir`.
89 Returns
90 -------
91 context : `contextlib.ContextManager`
92 A context manager that returns the new directory name on ``__enter__``
93 and removes the temporary directory (via `removeTestTempDir`) on
94 ``__exit__``.
95 """
96 root = makeTestTempDir(default_base)
97 try:
98 yield root
99 finally:
100 removeTestTempDir(root)
103class ButlerTestHelper:
104 """Mixin with helpers for unit tests."""
106 def assertAstropyTablesEqual(self, tables, expectedTables, filterColumns=False):
107 """Verify that a list of astropy tables matches a list of expected
108 astropy tables.
110 Parameters
111 ----------
112 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
113 The table or tables that should match the expected tables.
114 expectedTables : `astropy.table.Table`
115 or iterable [`astropy.table.Table`]
116 The tables with expected values to which the tables under test will
117 be compared.
118 filterColumns : `bool`
119 If `True` then only compare columns that exist in
120 ``expectedTables``.
121 """
122 # If a single table is passed in for tables or expectedTables, put it
123 # in a list.
124 if isinstance(tables, AstropyTable):
125 tables = [tables]
126 if isinstance(expectedTables, AstropyTable):
127 expectedTables = [expectedTables]
128 diff = io.StringIO()
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 # Assert that they match:
138 self.assertTrue(report_diff_values(table, expected, fileobj=diff), msg="\n" + diff.getvalue())
141def readTable(textTable):
142 """Read an astropy table from formatted text.
144 Contains formatting that causes the astropy table to print an empty string
145 instead of "--" for missing/unpopulated values in the text table.
148 Parameters
149 ----------
150 textTable : `str`
151 The text version of the table to read.
153 Returns
154 -------
155 table : `astropy.table.Table`
156 The table as an astropy table.
157 """
158 return AstropyTable.read(
159 textTable,
160 format="ascii",
161 data_start=2, # skip the header row and the header row underlines.
162 fill_values=[("", 0, "")],
163 )
166class MetricTestRepo:
167 """Creates and manage a test repository on disk with datasets that
168 may be queried and modified for unit tests.
170 Parameters
171 ----------
172 root : `str`
173 The location of the repository, to pass to ``Butler.makeRepo``.
174 configFile : `str`
175 The path to the config file, to pass to ``Butler.makeRepo``.
176 """
178 @staticmethod
179 def _makeExampleMetrics():
180 """Make an object to put into the repository."""
181 return MetricsExample(
182 {"AM1": 5.2, "AM2": 30.6},
183 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
184 [563, 234, 456.7, 752, 8, 9, 27],
185 )
187 @staticmethod
188 def _makeDimensionData(id, name, datetimeBegin=None, datetimeEnd=None):
189 """Make a dict of dimensional data with default values to insert into
190 the registry.
191 """
192 data = dict(instrument="DummyCamComp", id=id, name=name, physical_filter="d-r", visit_system=1)
193 if datetimeBegin:
194 data["datetime_begin"] = datetimeBegin
195 data["datetime_end"] = datetimeEnd
196 return data
198 def __init__(self, root, configFile):
199 self.root = root
200 Butler.makeRepo(self.root, config=Config(configFile))
201 butlerConfigFile = os.path.join(self.root, "butler.yaml")
202 self.storageClassFactory = StorageClassFactory()
203 self.storageClassFactory.addFromConfig(butlerConfigFile)
205 # New datasets will be added to run and tag, but we will only look in
206 # tag when looking up datasets.
207 run = "ingest/run"
208 tag = "ingest"
209 self.butler = Butler(butlerConfigFile, run=run, collections=[tag])
210 self.butler.registry.registerCollection(tag, CollectionType.TAGGED)
212 # Create and register a DatasetType
213 self.datasetType = addDatasetType(
214 self.butler, "test_metric_comp", ("instrument", "visit"), "StructuredCompositeReadComp"
215 )
217 # Add needed Dimensions
218 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
219 self.butler.registry.insertDimensionData(
220 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"}
221 )
222 self.butler.registry.insertDimensionData(
223 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"}
224 )
225 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
226 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
227 self.butler.registry.insertDimensionData(
228 "visit",
229 dict(
230 instrument="DummyCamComp",
231 id=423,
232 name="fourtwentythree",
233 physical_filter="d-r",
234 visit_system=1,
235 datetimeBegin=visitStart,
236 datetimeEnd=visitEnd,
237 ),
238 )
239 self.butler.registry.insertDimensionData(
240 "visit",
241 dict(
242 instrument="DummyCamComp",
243 id=424,
244 name="fourtwentyfour",
245 physical_filter="d-r",
246 visit_system=1,
247 ),
248 )
250 self.addDataset({"instrument": "DummyCamComp", "visit": 423})
251 self.addDataset({"instrument": "DummyCamComp", "visit": 424})
253 def addDataset(self, dataId, run=None, datasetType=None):
254 """Create a new example metric and add it to the named run with the
255 given dataId.
257 Overwrites tags, so this does not try to associate the new dataset with
258 existing tags. (If/when tags are needed this can be added to the
259 arguments of this function.)
261 Parameters
262 ----------
263 dataId : `dict`
264 The dataId for the new metric.
265 run : `str`, optional
266 The name of the run to create and add a dataset to. If `None`, the
267 dataset will be added to the root butler.
268 datasetType : ``DatasetType``, optional
269 The dataset type of the added dataset. If `None`, will use the
270 default dataset type.
271 """
272 if run:
273 self.butler.registry.registerCollection(run, type=CollectionType.RUN)
274 metric = self._makeExampleMetrics()
275 self.butler.put(metric, self.datasetType if datasetType is None else datasetType, dataId, run=run)