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

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 

26import io 

27import os 

28import shutil 

29import tempfile 

30from contextlib import contextmanager 

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 Butler, Config, StorageClassFactory 

38from ..registry import CollectionType 

39from ..tests import MetricsExample, addDatasetType 

40 

41 

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

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

44 

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

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

47 not. 

48 

49 Parameters 

50 ---------- 

51 default_base : `str` 

52 Default parent directory. 

53 

54 Returns 

55 ------- 

56 dir : `str` 

57 Name of the new temporary directory. 

58 """ 

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

60 return tempfile.mkdtemp(dir=base) 

61 

62 

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

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

65 unable to. 

66 

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

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

69 

70 Parameters 

71 ---------- 

72 root : `str`, optional 

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

74 """ 

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

76 shutil.rmtree(root, ignore_errors=True) 

77 

78 

79@contextmanager 

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

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

82 attempts to remove it. 

83 

84 Parameters 

85 ---------- 

86 default_base : `str` 

87 Default parent directory, forwarded to `makeTestTempDir`. 

88 

89 Returns 

90 ------- 

91 context : `contextlib.ContextManager` 

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

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

94 ``__exit__``. 

95 """ 

96 root = makeTestTempDir(default_base) 

97 try: 

98 yield root 

99 finally: 

100 removeTestTempDir(root) 

101 

102 

103class ButlerTestHelper: 

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

105 

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

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

108 astropy tables. 

109 

110 Parameters 

111 ---------- 

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

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

114 expectedTables : `astropy.table.Table` 

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

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

117 be compared. 

118 filterColumns : `bool` 

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

120 ``expectedTables``. 

121 """ 

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

123 # in a list. 

124 if isinstance(tables, AstropyTable): 

125 tables = [tables] 

126 if isinstance(expectedTables, AstropyTable): 

127 expectedTables = [expectedTables] 

128 diff = io.StringIO() 

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 # Assert that they match: 

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

139 

140 

141def readTable(textTable): 

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

143 

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

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

146 

147 

148 Parameters 

149 ---------- 

150 textTable : `str` 

151 The text version of the table to read. 

152 

153 Returns 

154 ------- 

155 table : `astropy.table.Table` 

156 The table as an astropy table. 

157 """ 

158 return AstropyTable.read( 

159 textTable, 

160 format="ascii", 

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

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

163 ) 

164 

165 

166class MetricTestRepo: 

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

168 may be queried and modified for unit tests. 

169 

170 Parameters 

171 ---------- 

172 root : `str` 

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

174 configFile : `str` 

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

176 """ 

177 

178 @staticmethod 

179 def _makeExampleMetrics(): 

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

181 return MetricsExample( 

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

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

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

185 ) 

186 

187 @staticmethod 

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

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

190 the registry. 

191 """ 

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

193 if datetimeBegin: 

194 data["datetime_begin"] = datetimeBegin 

195 data["datetime_end"] = datetimeEnd 

196 return data 

197 

198 def __init__(self, root, configFile): 

199 self.root = root 

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

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

202 self.storageClassFactory = StorageClassFactory() 

203 self.storageClassFactory.addFromConfig(butlerConfigFile) 

204 

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

206 # tag when looking up datasets. 

207 run = "ingest/run" 

208 tag = "ingest" 

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

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

211 

212 # Create and register a DatasetType 

213 self.datasetType = addDatasetType( 

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

215 ) 

216 

217 # Add needed Dimensions 

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

219 self.butler.registry.insertDimensionData( 

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

221 ) 

222 self.butler.registry.insertDimensionData( 

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

224 ) 

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

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

227 self.butler.registry.insertDimensionData( 

228 "visit", 

229 dict( 

230 instrument="DummyCamComp", 

231 id=423, 

232 name="fourtwentythree", 

233 physical_filter="d-r", 

234 visit_system=1, 

235 datetimeBegin=visitStart, 

236 datetimeEnd=visitEnd, 

237 ), 

238 ) 

239 self.butler.registry.insertDimensionData( 

240 "visit", 

241 dict( 

242 instrument="DummyCamComp", 

243 id=424, 

244 name="fourtwentyfour", 

245 physical_filter="d-r", 

246 visit_system=1, 

247 ), 

248 ) 

249 

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

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

252 

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

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

255 given dataId. 

256 

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

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

259 arguments of this function.) 

260 

261 Parameters 

262 ---------- 

263 dataId : `dict` 

264 The dataId for the new metric. 

265 run : `str`, optional 

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

267 dataset will be added to the root butler. 

268 datasetType : ``DatasetType``, optional 

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

270 default dataset type. 

271 """ 

272 if run: 

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

274 metric = self._makeExampleMetrics() 

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