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

86 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-27 09:44 +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/>. 

27 

28from __future__ import annotations 

29 

30__all__ = () 

31 

32import os 

33import shutil 

34import tempfile 

35from collections.abc import Callable, Iterator, Sequence 

36from contextlib import contextmanager 

37from typing import TYPE_CHECKING, Any 

38 

39import astropy 

40from astropy.table import Table as AstropyTable 

41 

42from .. import Butler, Config, StorageClassFactory 

43from ..registry import CollectionType 

44from ..tests import MetricsExample, addDatasetType 

45 

46if TYPE_CHECKING: 

47 import unittest 

48 

49 from lsst.daf.butler import DatasetType 

50 

51 class TestCaseMixin(unittest.TestCase): 

52 """Base class for mixin test classes that use TestCase methods.""" 

53 

54 pass 

55 

56else: 

57 

58 class TestCaseMixin: 

59 """Do-nothing definition of mixin base class for regular execution.""" 

60 

61 pass 

62 

63 

64def makeTestTempDir(default_base: str) -> str: 

65 """Create a temporary directory for test usage. 

66 

67 The directory will be created within ``DAF_BUTLER_TEST_TMP`` if that 

68 environment variable is set, falling back to ``default_base`` if it is 

69 not. 

70 

71 Parameters 

72 ---------- 

73 default_base : `str` 

74 Default parent directory. 

75 

76 Returns 

77 ------- 

78 dir : `str` 

79 Name of the new temporary directory. 

80 """ 

81 base = os.environ.get("DAF_BUTLER_TEST_TMP", default_base) 

82 return tempfile.mkdtemp(dir=base) 

83 

84 

85def removeTestTempDir(root: str | None) -> None: 

86 """Attempt to remove a temporary test directory, but do not raise if 

87 unable to. 

88 

89 Unlike `tempfile.TemporaryDirectory`, this passes ``ignore_errors=True`` 

90 to ``shutil.rmtree`` at close, making it safe to use on NFS. 

91 

92 Parameters 

93 ---------- 

94 root : `str`, optional 

95 Name of the directory to be removed. If `None`, nothing will be done. 

96 """ 

97 if root is not None and os.path.exists(root): 

98 shutil.rmtree(root, ignore_errors=True) 

99 

100 

101@contextmanager 

102def safeTestTempDir(default_base: str) -> Iterator[str]: 

103 """Return a context manager that creates a temporary directory and then 

104 attempts to remove it. 

105 

106 Parameters 

107 ---------- 

108 default_base : `str` 

109 Default parent directory, forwarded to `makeTestTempDir`. 

110 

111 Returns 

112 ------- 

113 context : `contextlib.ContextManager` 

114 A context manager that returns the new directory name on ``__enter__`` 

115 and removes the temporary directory (via `removeTestTempDir`) on 

116 ``__exit__``. 

117 """ 

118 root = makeTestTempDir(default_base) 

119 try: 

120 yield root 

121 finally: 

122 removeTestTempDir(root) 

123 

124 

125class ButlerTestHelper: 

126 """Mixin with helpers for unit tests.""" 

127 

128 assertEqual: Callable 

129 assertIsInstance: Callable 

130 maxDiff: int | None 

131 

132 def assertAstropyTablesEqual( 

133 self, 

134 tables: AstropyTable | Sequence[AstropyTable], 

135 expectedTables: AstropyTable | Sequence[AstropyTable], 

136 filterColumns: bool = False, 

137 unorderedRows: bool = False, 

138 ) -> None: 

139 """Verify that a list of astropy tables matches a list of expected 

140 astropy tables. 

141 

142 Parameters 

143 ---------- 

144 tables : `astropy.table.Table` or iterable [`astropy.table.Table`] 

145 The table or tables that should match the expected tables. 

146 expectedTables : `astropy.table.Table` 

147 or iterable [`astropy.table.Table`] 

148 The tables with expected values to which the tables under test will 

149 be compared. 

150 filterColumns : `bool` 

151 If `True` then only compare columns that exist in 

152 ``expectedTables``. 

153 unorderedRows : `bool`, optional 

154 If `True` (`False` is default), don't require tables to have their 

155 rows in the same order. 

156 """ 

157 # If a single table is passed in for tables or expectedTables, put it 

158 # in a list. 

159 if isinstance(tables, AstropyTable): 

160 tables = [tables] 

161 if isinstance(expectedTables, AstropyTable): 

162 expectedTables = [expectedTables] 

163 self.assertEqual(len(tables), len(expectedTables)) 

164 for table, expected in zip(tables, expectedTables, strict=True): 

165 # Assert that we are testing what we think we are testing: 

166 self.assertIsInstance(table, AstropyTable) 

167 self.assertIsInstance(expected, AstropyTable) 

168 if filterColumns: 

169 table = table.copy() 

170 table.keep_columns(expected.colnames) 

171 if unorderedRows: 

172 table = table.copy() 

173 table.sort(table.colnames) 

174 expected = expected.copy() 

175 expected.sort(expected.colnames) 

176 # Assert that they match. 

177 # Recommendation from Astropy Slack is to format the table into 

178 # lines for comparison. We do not compare column data types. 

179 table1 = table.pformat_all() 

180 expected1 = expected.pformat_all() 

181 original_max = self.maxDiff 

182 self.maxDiff = None # This is required to get the full diff. 

183 try: 

184 self.assertEqual(table1, expected1) 

185 finally: 

186 self.maxDiff = original_max 

187 

188 

189def readTable(textTable: str) -> AstropyTable: 

190 """Read an astropy table from formatted text. 

191 

192 Contains formatting that causes the astropy table to print an empty string 

193 instead of "--" for missing/unpopulated values in the text table. 

194 

195 

196 Parameters 

197 ---------- 

198 textTable : `str` 

199 The text version of the table to read. 

200 

201 Returns 

202 ------- 

203 table : `astropy.table.Table` 

204 The table as an astropy table. 

205 """ 

206 return AstropyTable.read( 

207 textTable, 

208 format="ascii", 

209 data_start=2, # skip the header row and the header row underlines. 

210 fill_values=[("", 0, "")], 

211 ) 

212 

213 

214class MetricTestRepo: 

215 """Creates and manage a test repository on disk with datasets that 

216 may be queried and modified for unit tests. 

217 

218 Parameters 

219 ---------- 

220 root : `str` 

221 The location of the repository, to pass to ``Butler.makeRepo``. 

222 configFile : `str` 

223 The path to the config file, to pass to ``Butler.makeRepo``. 

224 """ 

225 

226 @staticmethod 

227 def _makeExampleMetrics() -> MetricsExample: 

228 """Make an object to put into the repository.""" 

229 return MetricsExample( 

230 {"AM1": 5.2, "AM2": 30.6}, 

231 {"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}}, 

232 [563, 234, 456.7, 752, 8, 9, 27], 

233 ) 

234 

235 def __init__(self, root: str, configFile: str) -> None: 

236 self.root = root 

237 Butler.makeRepo(self.root, config=Config(configFile)) 

238 butlerConfigFile = os.path.join(self.root, "butler.yaml") 

239 self.storageClassFactory = StorageClassFactory() 

240 self.storageClassFactory.addFromConfig(butlerConfigFile) 

241 

242 # New datasets will be added to run and tag, but we will only look in 

243 # tag when looking up datasets. 

244 run = "ingest/run" 

245 tag = "ingest" 

246 self.butler = Butler.from_config(butlerConfigFile, run=run, collections=[tag]) 

247 self.butler.registry.registerCollection(tag, CollectionType.TAGGED) 

248 

249 # Create and register a DatasetType 

250 self.datasetType = addDatasetType( 

251 self.butler, "test_metric_comp", {"instrument", "visit"}, "StructuredCompositeReadComp" 

252 ) 

253 

254 # Add needed Dimensions 

255 self.butler.registry.insertDimensionData("instrument", {"name": "DummyCamComp"}) 

256 self.butler.registry.insertDimensionData( 

257 "physical_filter", {"instrument": "DummyCamComp", "name": "d-r", "band": "R"} 

258 ) 

259 self.butler.registry.insertDimensionData( 

260 "visit_system", {"instrument": "DummyCamComp", "id": 1, "name": "default"} 

261 ) 

262 visitStart = astropy.time.Time("2020-01-01 08:00:00.123456789", scale="tai") 

263 visitEnd = astropy.time.Time("2020-01-01 08:00:36.66", scale="tai") 

264 self.butler.registry.insertDimensionData( 

265 "visit", 

266 dict( 

267 instrument="DummyCamComp", 

268 id=423, 

269 name="fourtwentythree", 

270 physical_filter="d-r", 

271 datetimeBegin=visitStart, 

272 datetimeEnd=visitEnd, 

273 ), 

274 ) 

275 self.butler.registry.insertDimensionData( 

276 "visit", 

277 dict( 

278 instrument="DummyCamComp", 

279 id=424, 

280 name="fourtwentyfour", 

281 physical_filter="d-r", 

282 ), 

283 ) 

284 

285 self.addDataset({"instrument": "DummyCamComp", "visit": 423}) 

286 self.addDataset({"instrument": "DummyCamComp", "visit": 424}) 

287 

288 def addDataset( 

289 self, dataId: dict[str, Any], run: str | None = None, datasetType: DatasetType | None = None 

290 ) -> None: 

291 """Create a new example metric and add it to the named run with the 

292 given dataId. 

293 

294 Overwrites tags, so this does not try to associate the new dataset with 

295 existing tags. (If/when tags are needed this can be added to the 

296 arguments of this function.) 

297 

298 Parameters 

299 ---------- 

300 dataId : `dict` 

301 The dataId for the new metric. 

302 run : `str`, optional 

303 The name of the run to create and add a dataset to. If `None`, the 

304 dataset will be added to the root butler. 

305 datasetType : ``DatasetType``, optional 

306 The dataset type of the added dataset. If `None`, will use the 

307 default dataset type. 

308 """ 

309 if run: 

310 self.butler.registry.registerCollection(run, type=CollectionType.RUN) 

311 metric = self._makeExampleMetrics() 

312 self.butler.put(metric, self.datasetType if datasetType is None else datasetType, dataId, run=run)