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

86 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-12 10:56 -0700

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/>. 

21 

22from __future__ import annotations 

23 

24__all__ = () 

25 

26import os 

27import shutil 

28import tempfile 

29from collections.abc import Callable, Iterator, Sequence 

30from contextlib import contextmanager 

31from typing import TYPE_CHECKING, Any 

32 

33import astropy 

34from astropy.table import Table as AstropyTable 

35 

36from .. import Butler, Config, StorageClassFactory 

37from ..registry import CollectionType 

38from ..tests import MetricsExample, addDatasetType 

39 

40if TYPE_CHECKING: 

41 import unittest 

42 

43 from lsst.daf.butler import DatasetType 

44 

45 class TestCaseMixin(unittest.TestCase): 

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

47 

48 pass 

49 

50else: 

51 

52 class TestCaseMixin: 

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

54 

55 pass 

56 

57 

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

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

60 

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. 

64 

65 Parameters 

66 ---------- 

67 default_base : `str` 

68 Default parent directory. 

69 

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) 

77 

78 

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

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

81 unable to. 

82 

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

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

85 

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) 

93 

94 

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. 

99 

100 Parameters 

101 ---------- 

102 default_base : `str` 

103 Default parent directory, forwarded to `makeTestTempDir`. 

104 

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) 

117 

118 

119class ButlerTestHelper: 

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

121 

122 assertEqual: Callable 

123 assertIsInstance: Callable 

124 maxDiff: int | None 

125 

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. 

135 

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 

181 

182 

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

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

185 

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

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

188 

189 

190 Parameters 

191 ---------- 

192 textTable : `str` 

193 The text version of the table to read. 

194 

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 ) 

206 

207 

208class MetricTestRepo: 

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

210 may be queried and modified for unit tests. 

211 

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 """ 

219 

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 ) 

228 

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) 

235 

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) 

242 

243 # Create and register a DatasetType 

244 self.datasetType = addDatasetType( 

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

246 ) 

247 

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 ) 

278 

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

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

281 

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. 

287 

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.) 

291 

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)