Hide keyboard shortcuts

Hot-keys 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

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 

22"""Unit tests for daf_butler CLI query-collections command. 

23""" 

24 

25import astropy 

26from astropy.table import Table as AstropyTable 

27from astropy.utils.diff import report_diff_values 

28from numpy import array 

29import io 

30import os 

31import shutil 

32import tempfile 

33import unittest 

34 

35from lsst.daf.butler import ( 

36 Butler, 

37 Config, 

38 DatasetRef, 

39 DatasetType, 

40 StorageClassFactory 

41) 

42from lsst.daf.butler import script 

43from lsst.daf.butler.tests import MetricsExample 

44 

45 

46TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

47 

48 

49class QueryDatasetsTest(unittest.TestCase): 

50 

51 configFile = os.path.join(TESTDIR, "config/basic/butler.yaml") 

52 storageClassFactory = StorageClassFactory() 

53 

54 def _assertTablesEqual(self, tables, expectedTables): 

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

56 astropy tables.""" 

57 diff = io.StringIO() 

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

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

60 self.assertTrue(report_diff_values(table, expected, fileobj=diff), msg=diff.getvalue()) 

61 

62 @staticmethod 

63 def _makeExampleMetrics(): 

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

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

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

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

68 

69 @staticmethod 

70 def _addDatasetType(datasetTypeName, dimensions, storageClass, registry): 

71 """Create a DatasetType and register it 

72 """ 

73 datasetType = DatasetType(datasetTypeName, dimensions, storageClass) 

74 registry.registerDatasetType(datasetType) 

75 return datasetType 

76 

77 @staticmethod 

78 def _queryDatasets(repo, glob=(), collections=(), where="", find_first=False, show_uri=False): 

79 return script.queryDatasets(repo, glob, collections, where, find_first, show_uri) 

80 

81 def setUp(self): 

82 self.root = tempfile.mkdtemp(dir=TESTDIR) 

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

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

85 self.storageClassFactory.addFromConfig(self.configFile) 

86 

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

88 # tag when looking up datasets. 

89 run = "ingest/run" 

90 tag = "ingest" 

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

92 

93 # There will not be a collection yet 

94 collections = set(self.butler.registry.queryCollections()) 

95 self.assertEqual(collections, set([run, tag])) 

96 

97 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp") 

98 

99 # Create and register a DatasetType 

100 dimensions = self.butler.registry.dimensions.extract(["instrument", "visit"]) 

101 datasetTypeName = "test_metric_comp" 

102 self.datasetType = self._addDatasetType(datasetTypeName, dimensions, storageClass, 

103 self.butler.registry) 

104 

105 # Add needed Dimensions 

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

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

108 "name": "d-r", 

109 "band": "R"}) 

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

111 "id": 1, 

112 "name": "default"}) 

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

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

115 self.butler.registry.insertDimensionData("visit", 

116 {"instrument": "DummyCamComp", "id": 423, 

117 "name": "fourtwentythree", "physical_filter": "d-r", 

118 "visit_system": 1, "datetime_begin": visit_start, 

119 "datetime_end": visit_end}) 

120 self.butler.registry.insertDimensionData("visit", {"instrument": "DummyCamComp", "id": 424, 

121 "name": "fourtwentyfour", "physical_filter": "d-r", 

122 "visit_system": 1}) 

123 metric = self._makeExampleMetrics() 

124 dataId = {"instrument": "DummyCamComp", "visit": 423} 

125 ref = DatasetRef(self.datasetType, dataId, id=None) 

126 self.butler.put(metric, ref) 

127 

128 metric = self._makeExampleMetrics() 

129 dataId = {"instrument": "DummyCamComp", "visit": 424} 

130 ref = DatasetRef(self.datasetType, dataId, id=None) 

131 self.butler.put(metric, ref) 

132 

133 def tearDown(self): 

134 if os.path.exists(self.root): 

135 shutil.rmtree(self.root, ignore_errors=True) 

136 

137 def testShowURI(self): 

138 """Test for expected output with show_uri=True.""" 

139 tables = self._queryDatasets(repo=self.butlerConfigFile, show_uri=True) 

140 

141 expectedTables = ( 

142 AstropyTable(array(( 

143 ("test_metric_comp.data", "ingest/run", "1", "DummyCamComp", "423", 

144 self.butler.datastore.root.join( 

145 "ingest/run/test_metric_comp.data/" 

146 "test_metric_comp_v00000423_fDummyCamComp_data.yaml")), 

147 ("test_metric_comp.data", "ingest/run", "2", "DummyCamComp", "424", 

148 self.butler.datastore.root.join( 

149 "ingest/run/test_metric_comp.data/" 

150 "test_metric_comp_v00000424_fDummyCamComp_data.yaml")))), 

151 names=("type", "run", "id", "instrument", "visit", "URI")), 

152 AstropyTable(array(( 

153 ("test_metric_comp.output", "ingest/run", "1", "DummyCamComp", "423", 

154 self.butler.datastore.root.join( 

155 "ingest/run/test_metric_comp.output/" 

156 "test_metric_comp_v00000423_fDummyCamComp_output.yaml")), 

157 ("test_metric_comp.output", "ingest/run", "2", "DummyCamComp", "424", 

158 self.butler.datastore.root.join( 

159 "ingest/run/test_metric_comp.output/" 

160 "test_metric_comp_v00000424_fDummyCamComp_output.yaml")))), 

161 names=("type", "run", "id", "instrument", "visit", "URI")), 

162 AstropyTable(array(( 

163 ("test_metric_comp.summary", "ingest/run", "1", "DummyCamComp", "423", 

164 self.butler.datastore.root.join( 

165 "ingest/run/test_metric_comp.summary/" 

166 "test_metric_comp_v00000423_fDummyCamComp_summary.yaml")), 

167 ("test_metric_comp.summary", "ingest/run", "2", "DummyCamComp", "424", 

168 self.butler.datastore.root.join( 

169 "ingest/run/test_metric_comp.summary/" 

170 "test_metric_comp_v00000424_fDummyCamComp_summary.yaml")))), 

171 names=("type", "run", "id", "instrument", "visit", "URI")) 

172 ) 

173 

174 self._assertTablesEqual(tables, expectedTables) 

175 

176 def testNoShowURI(self): 

177 """Test for expected output without show_uri (default is False).""" 

178 tables = self._queryDatasets(repo=self.butlerConfigFile) 

179 

180 expectedTables = ( 

181 AstropyTable(array(( 

182 ("test_metric_comp", "ingest/run", "1", "DummyCamComp", "423"), 

183 ("test_metric_comp", "ingest/run", "2", "DummyCamComp", "424"))), 

184 names=("type", "run", "id", "instrument", "visit") 

185 ), 

186 ) 

187 

188 self._assertTablesEqual(tables, expectedTables) 

189 

190 def testWhere(self): 

191 """Test using the where clause to reduce the number of rows returned. 

192 """ 

193 tables = self._queryDatasets(repo=self.butlerConfigFile, where="visit=423") 

194 

195 expectedTables = ( 

196 AstropyTable(array( 

197 ("test_metric_comp", "ingest/run", "1", "DummyCamComp", "423")), 

198 names=("type", "run", "id", "instrument", "visit") 

199 ) 

200 ) 

201 

202 self._assertTablesEqual(tables, expectedTables) 

203 

204 def testGlobDatasetType(self): 

205 """Test specifying dataset type.""" 

206 # Create and register an additional DatasetType 

207 dimensions = self.butler.registry.dimensions.extract(["instrument", "visit"]) 

208 datasetTypeName = "alt_test_metric_comp" 

209 storageClass = self.storageClassFactory.getStorageClass("StructuredCompositeReadComp") 

210 datasetType = self._addDatasetType(datasetTypeName, dimensions, storageClass, self.butler.registry) 

211 self.butler.registry.insertDimensionData( 

212 "visit", {"instrument": "DummyCamComp", "id": 425, 

213 "name": "fourtwentyfive", "physical_filter": "d-r", 

214 "visit_system": 1}) 

215 metric = self._makeExampleMetrics() 

216 dataId = {"instrument": "DummyCamComp", "visit": 425} 

217 ref = DatasetRef(datasetType, dataId, id=None) 

218 self.butler.put(metric, ref) 

219 

220 # verify the new dataset type increases the number of tables found: 

221 tables = self._queryDatasets(repo=self.butlerConfigFile) 

222 

223 expectedTables = ( 

224 AstropyTable(array(( 

225 ("test_metric_comp", "ingest/run", "1", "DummyCamComp", "423"), 

226 ("test_metric_comp", "ingest/run", "2", "DummyCamComp", "424"))), 

227 names=("type", "run", "id", "instrument", "visit") 

228 ), 

229 AstropyTable(array(( 

230 ("alt_test_metric_comp", "ingest/run", "3", "DummyCamComp", "425"))), 

231 names=("type", "run", "id", "instrument", "visit") 

232 ) 

233 ) 

234 

235 self._assertTablesEqual(tables, expectedTables) 

236 

237 def testFindFirstAndCollections(self): 

238 """Test the find-first option, and the collections option, since it 

239 is required for find-first.""" 

240 

241 # Create a new butler with a new run, and add a dataset to shadow an 

242 # existing dataset. 

243 self.butler = Butler(self.root, run="foo") 

244 metric = self._makeExampleMetrics() 

245 dataId = {"instrument": "DummyCamComp", "visit": 424} 

246 ref = DatasetRef(self.datasetType, dataId, id=None) 

247 self.butler.put(metric, ref) 

248 

249 # Verify that without find-first, duplicate datasets are returned 

250 tables = self._queryDatasets(repo=self.butlerConfigFile, 

251 collections=["foo", "ingest/run"], 

252 show_uri=True) 

253 

254 expectedTables = ( 

255 AstropyTable(array( 

256 ( 

257 ("test_metric_comp.data", "foo", "3", "DummyCamComp", "424", 

258 self.butler.datastore.root.join( 

259 "foo/test_metric_comp.data/" 

260 "test_metric_comp_v00000424_fDummyCamComp_data.yaml")), 

261 ("test_metric_comp.data", "ingest/run", "1", "DummyCamComp", "423", 

262 self.butler.datastore.root.join( 

263 "ingest/run/test_metric_comp.data/" 

264 "test_metric_comp_v00000423_fDummyCamComp_data.yaml")), 

265 ("test_metric_comp.data", "ingest/run", "2", "DummyCamComp", "424", 

266 self.butler.datastore.root.join( 

267 "ingest/run/test_metric_comp.data/" 

268 "test_metric_comp_v00000424_fDummyCamComp_data.yaml")), 

269 )), 

270 names=("type", "run", "id", "instrument", "visit", "URI")), 

271 AstropyTable(array( 

272 ( 

273 ("test_metric_comp.output", "foo", "3", "DummyCamComp", "424", 

274 self.butler.datastore.root.join( 

275 "foo/test_metric_comp.output/" 

276 "test_metric_comp_v00000424_fDummyCamComp_output.yaml")), 

277 ("test_metric_comp.output", "ingest/run", "1", "DummyCamComp", "423", 

278 self.butler.datastore.root.join( 

279 "ingest/run/test_metric_comp.output/" 

280 "test_metric_comp_v00000423_fDummyCamComp_output.yaml")), 

281 ("test_metric_comp.output", "ingest/run", "2", "DummyCamComp", "424", 

282 self.butler.datastore.root.join( 

283 "ingest/run/test_metric_comp.output/" 

284 "test_metric_comp_v00000424_fDummyCamComp_output.yaml")), 

285 )), 

286 names=("type", "run", "id", "instrument", "visit", "URI")), 

287 AstropyTable(array( 

288 ( 

289 ("test_metric_comp.summary", "foo", "3", "DummyCamComp", "424", 

290 self.butler.datastore.root.join( 

291 "foo/test_metric_comp.summary/" 

292 "test_metric_comp_v00000424_fDummyCamComp_summary.yaml")), 

293 ("test_metric_comp.summary", "ingest/run", "1", "DummyCamComp", "423", 

294 self.butler.datastore.root.join( 

295 "ingest/run/test_metric_comp.summary/" 

296 "test_metric_comp_v00000423_fDummyCamComp_summary.yaml")), 

297 ("test_metric_comp.summary", "ingest/run", "2", "DummyCamComp", "424", 

298 self.butler.datastore.root.join( 

299 "ingest/run/test_metric_comp.summary/" 

300 "test_metric_comp_v00000424_fDummyCamComp_summary.yaml")), 

301 )), 

302 names=("type", "run", "id", "instrument", "visit", "URI")) 

303 ) 

304 

305 self._assertTablesEqual(tables, expectedTables) 

306 

307 # Verify that with find first the duplicate dataset is eliminated and 

308 # the more recent dataset is returned. 

309 tables = self._queryDatasets(repo=self.butlerConfigFile, 

310 collections=["foo", "ingest/run"], 

311 show_uri=True, 

312 find_first=True) 

313 

314 expectedTables = ( 

315 AstropyTable(array( 

316 ( 

317 ("test_metric_comp.data", "foo", "3", "DummyCamComp", "424", 

318 self.butler.datastore.root.join( 

319 "foo/test_metric_comp.data/test_metric_comp_v00000424_fDummyCamComp_data.yaml")), 

320 ("test_metric_comp.data", "ingest/run", "1", "DummyCamComp", "423", 

321 self.butler.datastore.root.join( 

322 "ingest/run/test_metric_comp.data/" 

323 "test_metric_comp_v00000423_fDummyCamComp_data.yaml")), 

324 )), 

325 names=("type", "run", "id", "instrument", "visit", "URI")), 

326 AstropyTable(array( 

327 ( 

328 ("test_metric_comp.output", "foo", "3", "DummyCamComp", "424", 

329 self.butler.datastore.root.join( 

330 "foo/test_metric_comp.output/" 

331 "test_metric_comp_v00000424_fDummyCamComp_output.yaml")), 

332 ("test_metric_comp.output", "ingest/run", "1", "DummyCamComp", "423", 

333 self.butler.datastore.root.join( 

334 "ingest/run/test_metric_comp.output/" 

335 "test_metric_comp_v00000423_fDummyCamComp_output.yaml")), 

336 )), 

337 names=("type", "run", "id", "instrument", "visit", "URI")), 

338 AstropyTable(array( 

339 ( 

340 ("test_metric_comp.summary", "foo", "3", "DummyCamComp", "424", 

341 self.butler.datastore.root.join( 

342 "foo/test_metric_comp.summary/" 

343 "test_metric_comp_v00000424_fDummyCamComp_summary.yaml")), 

344 ("test_metric_comp.summary", "ingest/run", "1", "DummyCamComp", "423", 

345 self.butler.datastore.root.join( 

346 "ingest/run/test_metric_comp.summary/" 

347 "test_metric_comp_v00000423_fDummyCamComp_summary.yaml")), 

348 )), 

349 names=("type", "run", "id", "instrument", "visit", "URI")) 

350 ) 

351 

352 self._assertTablesEqual(tables, expectedTables) 

353 

354 

355if __name__ == "__main__": 355 ↛ 356line 355 didn't jump to line 356, because the condition on line 355 was never true

356 unittest.main()