Coverage for python/lsst/daf/butler/tests/utils.py : 32%

Hot-keys 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__ = ()
26from contextlib import contextmanager
27import io
28import os
29import shutil
30import tempfile
31from typing import Optional
33import astropy
34from astropy.table import Table as AstropyTable
35from astropy.utils.diff import report_diff_values
37from .. import (
38 Butler,
39 Config,
40 StorageClassFactory,
41)
42from ..tests import addDatasetType, MetricsExample
43from ..registry import CollectionType
46def makeTestTempDir(default_base: str) -> str:
47 """Create a temporary directory for test usage.
49 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that
50 environment variable is set, falling back to ``default_base`` if it is
51 not.
53 Parameters
54 ----------
55 default_base : `str`
56 Default parent directory.
58 Returns
59 -------
60 dir : `str`
61 Name of the new temporary directory.
62 """
63 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base)
64 return tempfile.mkdtemp(dir=base)
67def removeTestTempDir(root: Optional[str]) -> None:
68 """Attempt to remove a temporary test directory, but do not raise if
69 unable to.
71 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True``
72 to ``shutil.rmtree`` at close, making it safe to use on NFS.
74 Parameters
75 ----------
76 root : `str`, optional
77 Name of the directory to be removed. If `None`, nothing will be done.
78 """
79 if root is not None and os.path.exists(root):
80 shutil.rmtree(root, ignore_errors=True)
83@contextmanager
84def safeTestTempDir(default_base: str) -> str:
85 """Return a context manager that creates a temporary directory and then
86 attempts to remove it.
88 Parameters
89 ----------
90 default_base : `str`
91 Default parent directory, forwarded to `makeTestTempDir`.
93 Returns
94 -------
95 context : `contextlib.ContextManager`
96 A context manager that returns the new directory name on ``__enter__``
97 and removes the temporary directory (via `removeTestTempDir`) on
98 ``__exit__``.
99 """
100 root = makeTestTempDir(default_base)
101 try:
102 yield root
103 finally:
104 removeTestTempDir(root)
107class ButlerTestHelper:
108 """Mixin with helpers for unit tests."""
110 def assertAstropyTablesEqual(self, tables, expectedTables):
111 """Verify that a list of astropy tables matches a list of expected
112 astropy tables.
114 Parameters
115 ----------
116 tables : `astropy.table.Table` or iterable [`astropy.table.Table`]
117 The table or tables that should match the expected tables.
118 expectedTables : `astropy.table.Table`
119 or iterable [`astropy.table.Table`]
120 The tables with expected values to which the tables under test will
121 be compared.
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 diff = io.StringIO()
130 self.assertEqual(len(tables), len(expectedTables))
131 for table, expected in zip(tables, expectedTables):
132 # Assert that we are testing what we think we are testing:
133 self.assertIsInstance(table, AstropyTable)
134 self.assertIsInstance(expected, AstropyTable)
135 # Assert that they match:
136 self.assertTrue(report_diff_values(table, expected, fileobj=diff), msg="\n" + diff.getvalue())
139def readTable(textTable):
140 """Read an astropy table from formatted text.
142 Contains formatting that causes the astropy table to print an empty string
143 instead of "--" for missing/unpopulated values in the text table.
146 Parameters
147 ----------
148 textTable : `str`
149 The text version of the table to read.
151 Returns
152 -------
153 table : `astropy.table.Table`
154 The table as an astropy table.
155 """
156 return AstropyTable.read(textTable,
157 format="ascii",
158 data_start=2, # skip the header row and the header row underlines.
159 fill_values=[("", 0, "")])
162class MetricTestRepo:
163 """Creates and manage a test repository on disk with datasets that
164 may be queried and modified for unit tests.
166 Parameters
167 ----------
168 root : `str`
169 The location of the repository, to pass to ``Butler.makeRepo``.
170 configFile : `str`
171 The path to the config file, to pass to ``Butler.makeRepo``.
172 """
174 @staticmethod
175 def _makeExampleMetrics():
176 """Make an object to put into the repository.
177 """
178 return MetricsExample({"AM1": 5.2, "AM2": 30.6},
179 {"a": [1, 2, 3],
180 "b": {"blue": 5, "red": "green"}},
181 [563, 234, 456.7, 752, 8, 9, 27])
183 @staticmethod
184 def _makeDimensionData(id, name, datetimeBegin=None, datetimeEnd=None):
185 """Make a dict of dimensional data with default values to insert into
186 the registry.
187 """
188 data = dict(instrument="DummyCamComp",
189 id=id,
190 name=name,
191 physical_filter="d-r",
192 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(self.butler, "test_metric_comp", ("instrument", "visit"),
214 "StructuredCompositeReadComp")
216 # Add needed Dimensions
217 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"})
218 self.butler.registry.insertDimensionData("physical_filter", {"instrument": "DummyCamComp",
219 "name": "d-r",
220 "band": "R"})
221 self.butler.registry.insertDimensionData("visit_system", {"instrument": "DummyCamComp",
222 "id": 1,
223 "name": "default"})
224 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai")
225 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai")
226 self.butler.registry.insertDimensionData("visit", dict(instrument="DummyCamComp",
227 id=423,
228 name="fourtwentythree",
229 physical_filter="d-r",
230 visit_system=1,
231 datetimeBegin=visitStart,
232 datetimeEnd=visitEnd))
233 self.butler.registry.insertDimensionData("visit", dict(instrument="DummyCamComp",
234 id=424,
235 name="fourtwentyfour",
236 physical_filter="d-r",
237 visit_system=1))
239 self.addDataset({"instrument": "DummyCamComp", "visit": 423})
240 self.addDataset({"instrument": "DummyCamComp", "visit": 424})
242 def addDataset(self, dataId, run=None, datasetType=None):
243 """Create a new example metric and add it to the named run with the
244 given dataId.
246 Overwrites tags, so this does not try to associate the new dataset with
247 existing tags. (If/when tags are needed this can be added to the
248 arguments of this function.)
250 Parameters
251 ----------
252 dataId : `dict`
253 The dataId for the new metric.
254 run : `str`, optional
255 The name of the run to create and add a dataset to. If `None`, the
256 dataset will be added to the root butler.
257 datasetType : ``DatasetType``, optional
258 The dataset type of the added dataset. If `None`, will use the
259 default dataset type.
260 """
261 if run:
262 self.butler.registry.registerCollection(run, type=CollectionType.RUN)
263 metric = self._makeExampleMetrics()
264 self.butler.put(metric,
265 self.datasetType if datasetType is None else datasetType,
266 dataId,
267 run=run,
268 tags=())