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

104 statements  

« prev     ^ index     » next       coverage.py v7.4.4, created at 2024-04-19 03:44 -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 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 

30from unittest.mock import patch 

31 

32__all__ = () 

33 

34import os 

35import shutil 

36import tempfile 

37from collections.abc import Callable, Iterator, Sequence 

38from contextlib import contextmanager 

39from typing import TYPE_CHECKING, Any 

40 

41import astropy 

42from astropy.table import Table as AstropyTable 

43 

44from .. import Butler, Config, DatasetRef, StorageClassFactory, Timespan 

45from ..registry import CollectionType 

46from ..tests import MetricsExample, addDatasetType 

47 

48if TYPE_CHECKING: 

49 import unittest 

50 

51 from lsst.daf.butler import DatasetType 

52 

53 class TestCaseMixin(unittest.TestCase): 

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

55 

56 pass 

57 

58else: 

59 

60 class TestCaseMixin: 

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

62 

63 pass 

64 

65 

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

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

68 

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

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

71 not. 

72 

73 Parameters 

74 ---------- 

75 default_base : `str` 

76 Default parent directory. 

77 

78 Returns 

79 ------- 

80 dir : `str` 

81 Name of the new temporary directory. 

82 """ 

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

84 return tempfile.mkdtemp(dir=base) 

85 

86 

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

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

89 unable to. 

90 

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

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

93 

94 Parameters 

95 ---------- 

96 root : `str`, optional 

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

98 """ 

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

100 shutil.rmtree(root, ignore_errors=True) 

101 

102 

103@contextmanager 

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

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

106 attempts to remove it. 

107 

108 Parameters 

109 ---------- 

110 default_base : `str` 

111 Default parent directory, forwarded to `makeTestTempDir`. 

112 

113 Returns 

114 ------- 

115 context : `contextlib.ContextManager` 

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

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

118 ``__exit__``. 

119 """ 

120 root = makeTestTempDir(default_base) 

121 try: 

122 yield root 

123 finally: 

124 removeTestTempDir(root) 

125 

126 

127class ButlerTestHelper: 

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

129 

130 assertEqual: Callable 

131 assertIsInstance: Callable 

132 maxDiff: int | None 

133 

134 def assertAstropyTablesEqual( 

135 self, 

136 tables: AstropyTable | Sequence[AstropyTable], 

137 expectedTables: AstropyTable | Sequence[AstropyTable], 

138 filterColumns: bool = False, 

139 unorderedRows: bool = False, 

140 ) -> None: 

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

142 astropy tables. 

143 

144 Parameters 

145 ---------- 

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

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

148 expectedTables : `astropy.table.Table` 

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

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

151 be compared. 

152 filterColumns : `bool` 

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

154 ``expectedTables``. 

155 unorderedRows : `bool`, optional 

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

157 rows in the same order. 

158 """ 

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

160 # in a list. 

161 if isinstance(tables, AstropyTable): 

162 tables = [tables] 

163 if isinstance(expectedTables, AstropyTable): 

164 expectedTables = [expectedTables] 

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

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

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

168 self.assertIsInstance(table, AstropyTable) 

169 self.assertIsInstance(expected, AstropyTable) 

170 if filterColumns: 

171 table = table.copy() 

172 table.keep_columns(expected.colnames) 

173 if unorderedRows: 

174 table = table.copy() 

175 table.sort(table.colnames) 

176 expected = expected.copy() 

177 expected.sort(expected.colnames) 

178 # Assert that they match. 

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

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

181 table1 = table.pformat_all() 

182 expected1 = expected.pformat_all() 

183 original_max = self.maxDiff 

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

185 try: 

186 self.assertEqual(table1, expected1) 

187 finally: 

188 self.maxDiff = original_max 

189 

190 

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

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

193 

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

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

196 

197 Parameters 

198 ---------- 

199 textTable : `str` 

200 The text version of the table to read. 

201 

202 Returns 

203 ------- 

204 table : `astropy.table.Table` 

205 The table as an astropy table. 

206 """ 

207 return AstropyTable.read( 

208 textTable, 

209 format="ascii", 

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

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

212 ) 

213 

214 

215class MetricTestRepo: 

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

217 may be queried and modified for unit tests. 

218 

219 Parameters 

220 ---------- 

221 root : `str` 

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

223 configFile : `str` 

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

225 forceConfigRoot : `bool`, optional 

226 If `False`, any values present in the supplied ``config`` that 

227 would normally be reset are not overridden and will appear 

228 directly in the output config. Passed to ``Butler.makeRepo``. 

229 """ 

230 

231 METRICS_EXAMPLE_SUMMARY = {"AM1": 5.2, "AM2": 30.6} 

232 """The summary data included in ``MetricsExample`` objects stored in the 

233 test repo 

234 """ 

235 

236 _DEFAULT_RUN = "ingest/run" 

237 _DEFAULT_TAG = "ingest" 

238 

239 @staticmethod 

240 def _makeExampleMetrics() -> MetricsExample: 

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

242 return MetricsExample( 

243 MetricTestRepo.METRICS_EXAMPLE_SUMMARY, 

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

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

246 ) 

247 

248 def __init__(self, root: str, configFile: str, forceConfigRoot: bool = True) -> None: 

249 self.root = root 

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

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

252 butler = Butler.from_config(butlerConfigFile, run=self._DEFAULT_RUN, collections=[self._DEFAULT_TAG]) 

253 self._do_init(butler, butlerConfigFile) 

254 

255 @classmethod 

256 def create_from_butler(cls, butler: Butler, butler_config_file: str) -> MetricTestRepo: 

257 """Create a MetricTestRepo from an existing Butler instance. 

258 

259 Parameters 

260 ---------- 

261 butler : `Butler` 

262 `Butler` instance used for setting up the repository. 

263 butler_config_file : `str` 

264 Path to the config file used to set up that Butler instance. 

265 

266 Returns 

267 ------- 

268 repo : `MetricTestRepo` 

269 New instance of `MetricTestRepo` using the provided `Butler` 

270 instance. 

271 """ 

272 self = cls.__new__(cls) 

273 butler = butler._clone(run=self._DEFAULT_RUN, collections=[self._DEFAULT_TAG]) 

274 self._do_init(butler, butler_config_file) 

275 return self 

276 

277 def _do_init(self, butler: Butler, butlerConfigFile: str) -> None: 

278 self.butler = butler 

279 self.storageClassFactory = StorageClassFactory() 

280 self.storageClassFactory.addFromConfig(butlerConfigFile) 

281 

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

283 # tag when looking up datasets. 

284 self.butler.registry.registerCollection(self._DEFAULT_TAG, CollectionType.TAGGED) 

285 

286 # Create and register a DatasetType 

287 self.datasetType = addDatasetType( 

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

289 ) 

290 

291 # Add needed Dimensions 

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

293 self.butler.registry.insertDimensionData( 

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

295 ) 

296 self.butler.registry.insertDimensionData("day_obs", {"instrument": "DummyCamComp", "id": 20200101}) 

297 self.butler.registry.insertDimensionData( 

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

299 ) 

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

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

302 self.butler.registry.insertDimensionData( 

303 "visit", 

304 dict( 

305 instrument="DummyCamComp", 

306 id=423, 

307 name="fourtwentythree", 

308 physical_filter="d-r", 

309 timespan=Timespan(visitStart, visitEnd), 

310 day_obs=20200101, 

311 ), 

312 ) 

313 self.butler.registry.insertDimensionData( 

314 "visit", 

315 dict( 

316 instrument="DummyCamComp", 

317 id=424, 

318 name="fourtwentyfour", 

319 physical_filter="d-r", 

320 day_obs=20200101, 

321 ), 

322 ) 

323 

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

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

326 

327 def addDataset( 

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

329 ) -> DatasetRef: 

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

331 given dataId. 

332 

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

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

335 arguments of this function.) 

336 

337 Parameters 

338 ---------- 

339 dataId : `dict` 

340 The dataId for the new metric. 

341 run : `str`, optional 

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

343 dataset will be added to the root butler. 

344 datasetType : ``DatasetType``, optional 

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

346 default dataset type. 

347 

348 Returns 

349 ------- 

350 datasetRef : `DatasetRef` 

351 A reference to the added dataset. 

352 """ 

353 if run: 

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

355 else: 

356 run = self._DEFAULT_RUN 

357 metric = self._makeExampleMetrics() 

358 return self.butler.put( 

359 metric, self.datasetType if datasetType is None else datasetType, dataId, run=run 

360 ) 

361 

362 

363@contextmanager 

364def mock_env(new_environment: dict[str, str]) -> Iterator[None]: 

365 """Context manager to clear the process environment variables, replace them 

366 with new values, and restore them at the end of the test. 

367 

368 Parameters 

369 ---------- 

370 new_environment : `dict`[`str`, `str`] 

371 New environment variable values. 

372 """ 

373 with patch.dict(os.environ, new_environment, clear=True): 

374 yield