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

Shortcuts 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

79 statements  

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 

26from contextlib import contextmanager 

27import io 

28import os 

29import shutil 

30import tempfile 

31from typing import Optional 

32 

33import astropy 

34from astropy.table import Table as AstropyTable 

35from astropy.utils.diff import report_diff_values 

36 

37from .. import ( 

38 Butler, 

39 Config, 

40 StorageClassFactory, 

41) 

42from ..tests import addDatasetType, MetricsExample 

43from ..registry import CollectionType 

44 

45 

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

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

48 

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. 

52 

53 Parameters 

54 ---------- 

55 default_base : `str` 

56 Default parent directory. 

57 

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) 

65 

66 

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

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

69 unable to. 

70 

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

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

73 

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) 

81 

82 

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. 

87 

88 Parameters 

89 ---------- 

90 default_base : `str` 

91 Default parent directory, forwarded to `makeTestTempDir`. 

92 

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) 

105 

106 

107class ButlerTestHelper: 

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

109 

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

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

112 astropy tables. 

113 

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 filterColumns : `bool` 

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

124 ``expectedTables``. 

125 """ 

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

127 # in a list. 

128 if isinstance(tables, AstropyTable): 

129 tables = [tables] 

130 if isinstance(expectedTables, AstropyTable): 

131 expectedTables = [expectedTables] 

132 diff = io.StringIO() 

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

134 for table, expected in zip(tables, expectedTables): 

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

136 self.assertIsInstance(table, AstropyTable) 

137 self.assertIsInstance(expected, AstropyTable) 

138 if filterColumns: 

139 table = table.copy() 

140 table.keep_columns(expected.colnames) 

141 # Assert that they match: 

142 self.assertTrue(report_diff_values(table, expected, fileobj=diff), msg="\n" + diff.getvalue()) 

143 

144 

145def readTable(textTable): 

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

147 

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

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

150 

151 

152 Parameters 

153 ---------- 

154 textTable : `str` 

155 The text version of the table to read. 

156 

157 Returns 

158 ------- 

159 table : `astropy.table.Table` 

160 The table as an astropy table. 

161 """ 

162 return AstropyTable.read(textTable, 

163 format="ascii", 

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

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

166 

167 

168class MetricTestRepo: 

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

170 may be queried and modified for unit tests. 

171 

172 Parameters 

173 ---------- 

174 root : `str` 

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

176 configFile : `str` 

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

178 """ 

179 

180 @staticmethod 

181 def _makeExampleMetrics(): 

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

183 """ 

184 return MetricsExample({"AM1": 5.2, "AM2": 30.6}, 

185 {"a": [1, 2, 3], 

186 "b": {"blue": 5, "red": "green"}}, 

187 [563, 234, 456.7, 752, 8, 9, 27]) 

188 

189 @staticmethod 

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

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

192 the registry. 

193 """ 

194 data = dict(instrument="DummyCamComp", 

195 id=id, 

196 name=name, 

197 physical_filter="d-r", 

198 visit_system=1) 

199 if datetimeBegin: 

200 data["datetime_begin"] = datetimeBegin 

201 data["datetime_end"] = datetimeEnd 

202 return data 

203 

204 def __init__(self, root, configFile): 

205 self.root = root 

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

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

208 self.storageClassFactory = StorageClassFactory() 

209 self.storageClassFactory.addFromConfig(butlerConfigFile) 

210 

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

212 # tag when looking up datasets. 

213 run = "ingest/run" 

214 tag = "ingest" 

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

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

217 

218 # Create and register a DatasetType 

219 self.datasetType = addDatasetType(self.butler, "test_metric_comp", ("instrument", "visit"), 

220 "StructuredCompositeReadComp") 

221 

222 # Add needed Dimensions 

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

224 self.butler.registry.insertDimensionData("physical_filter", {"instrument": "DummyCamComp", 

225 "name": "d-r", 

226 "band": "R"}) 

227 self.butler.registry.insertDimensionData("visit_system", {"instrument": "DummyCamComp", 

228 "id": 1, 

229 "name": "default"}) 

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

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

232 self.butler.registry.insertDimensionData("visit", dict(instrument="DummyCamComp", 

233 id=423, 

234 name="fourtwentythree", 

235 physical_filter="d-r", 

236 visit_system=1, 

237 datetimeBegin=visitStart, 

238 datetimeEnd=visitEnd)) 

239 self.butler.registry.insertDimensionData("visit", dict(instrument="DummyCamComp", 

240 id=424, 

241 name="fourtwentyfour", 

242 physical_filter="d-r", 

243 visit_system=1)) 

244 

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

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

247 

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

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

250 given dataId. 

251 

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

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

254 arguments of this function.) 

255 

256 Parameters 

257 ---------- 

258 dataId : `dict` 

259 The dataId for the new metric. 

260 run : `str`, optional 

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

262 dataset will be added to the root butler. 

263 datasetType : ``DatasetType``, optional 

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

265 default dataset type. 

266 """ 

267 if run: 

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

269 metric = self._makeExampleMetrics() 

270 self.butler.put(metric, 

271 self.datasetType if datasetType is None else datasetType, 

272 dataId, 

273 run=run, 

274 tags=())