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

87 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-11-17 02:01 -0800

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 contextlib import contextmanager 

30from typing import Optional 

31 

32import astropy 

33from astropy.table import Table as AstropyTable 

34 

35from .. import Butler, Config, StorageClassFactory 

36from ..registry import CollectionType 

37from ..tests import MetricsExample, addDatasetType 

38 

39 

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

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

42 

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

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

45 not. 

46 

47 Parameters 

48 ---------- 

49 default_base : `str` 

50 Default parent directory. 

51 

52 Returns 

53 ------- 

54 dir : `str` 

55 Name of the new temporary directory. 

56 """ 

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

58 return tempfile.mkdtemp(dir=base) 

59 

60 

61def removeTestTempDir(root: Optional[str]) -> None: 

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

63 unable to. 

64 

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

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

67 

68 Parameters 

69 ---------- 

70 root : `str`, optional 

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

72 """ 

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

74 shutil.rmtree(root, ignore_errors=True) 

75 

76 

77@contextmanager 

78def safeTestTempDir(default_base: str) -> str: 

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

80 attempts to remove it. 

81 

82 Parameters 

83 ---------- 

84 default_base : `str` 

85 Default parent directory, forwarded to `makeTestTempDir`. 

86 

87 Returns 

88 ------- 

89 context : `contextlib.ContextManager` 

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

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

92 ``__exit__``. 

93 """ 

94 root = makeTestTempDir(default_base) 

95 try: 

96 yield root 

97 finally: 

98 removeTestTempDir(root) 

99 

100 

101class ButlerTestHelper: 

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

103 

104 def assertAstropyTablesEqual(self, tables, expectedTables, filterColumns=False, unorderedRows=False): 

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

106 astropy tables. 

107 

108 Parameters 

109 ---------- 

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

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

112 expectedTables : `astropy.table.Table` 

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

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

115 be compared. 

116 filterColumns : `bool` 

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

118 ``expectedTables``. 

119 unorderedRows : `bool`, optional 

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

121 rows in the same order. 

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 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 if unorderedRows: 

138 table = table.copy() 

139 table.sort(table.colnames) 

140 expected = expected.copy() 

141 expected.sort(expected.colnames) 

142 # Assert that they match. 

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

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

145 table1 = table.pformat_all() 

146 expected1 = expected.pformat_all() 

147 original_max = self.maxDiff 

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

149 try: 

150 self.assertEqual(table1, expected1) 

151 finally: 

152 self.maxDiff = original_max 

153 

154 

155def readTable(textTable): 

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

157 

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

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

160 

161 

162 Parameters 

163 ---------- 

164 textTable : `str` 

165 The text version of the table to read. 

166 

167 Returns 

168 ------- 

169 table : `astropy.table.Table` 

170 The table as an astropy table. 

171 """ 

172 return AstropyTable.read( 

173 textTable, 

174 format="ascii", 

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

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

177 ) 

178 

179 

180class MetricTestRepo: 

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

182 may be queried and modified for unit tests. 

183 

184 Parameters 

185 ---------- 

186 root : `str` 

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

188 configFile : `str` 

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

190 """ 

191 

192 @staticmethod 

193 def _makeExampleMetrics(): 

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

195 return MetricsExample( 

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

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

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

199 ) 

200 

201 @staticmethod 

202 def _makeDimensionData(id, name, datetimeBegin=None, datetimeEnd=None): 

203 """Make a dict of dimensional data with default values to insert into 

204 the registry. 

205 """ 

206 data = dict(instrument="DummyCamComp", id=id, name=name, physical_filter="d-r", visit_system=1) 

207 if datetimeBegin: 

208 data["datetime_begin"] = datetimeBegin 

209 data["datetime_end"] = datetimeEnd 

210 return data 

211 

212 def __init__(self, root, configFile): 

213 self.root = root 

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

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

216 self.storageClassFactory = StorageClassFactory() 

217 self.storageClassFactory.addFromConfig(butlerConfigFile) 

218 

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

220 # tag when looking up datasets. 

221 run = "ingest/run" 

222 tag = "ingest" 

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

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

225 

226 # Create and register a DatasetType 

227 self.datasetType = addDatasetType( 

228 self.butler, "test_metric_comp", ("instrument", "visit"), "StructuredCompositeReadComp" 

229 ) 

230 

231 # Add needed Dimensions 

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

233 self.butler.registry.insertDimensionData( 

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

235 ) 

236 self.butler.registry.insertDimensionData( 

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

238 ) 

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

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

241 self.butler.registry.insertDimensionData( 

242 "visit", 

243 dict( 

244 instrument="DummyCamComp", 

245 id=423, 

246 name="fourtwentythree", 

247 physical_filter="d-r", 

248 datetimeBegin=visitStart, 

249 datetimeEnd=visitEnd, 

250 ), 

251 ) 

252 self.butler.registry.insertDimensionData( 

253 "visit", 

254 dict( 

255 instrument="DummyCamComp", 

256 id=424, 

257 name="fourtwentyfour", 

258 physical_filter="d-r", 

259 ), 

260 ) 

261 

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

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

264 

265 def addDataset(self, dataId, run=None, datasetType=None): 

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

267 given dataId. 

268 

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

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

271 arguments of this function.) 

272 

273 Parameters 

274 ---------- 

275 dataId : `dict` 

276 The dataId for the new metric. 

277 run : `str`, optional 

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

279 dataset will be added to the root butler. 

280 datasetType : ``DatasetType``, optional 

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

282 default dataset type. 

283 """ 

284 if run: 

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

286 metric = self._makeExampleMetrics() 

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