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 pipe_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 <https://www.gnu.org/licenses/>. 

21 

22"""Unit tests for `lsst.pipe.base.tests`, a library for testing 

23PipelineTask subclasses. 

24""" 

25 

26import shutil 

27import tempfile 

28import unittest 

29 

30import lsst.utils.tests 

31import lsst.pex.config 

32import lsst.daf.butler 

33import lsst.daf.butler.tests as butlerTests 

34 

35from lsst.pipe.base import Struct, PipelineTask, PipelineTaskConfig, PipelineTaskConnections, connectionTypes 

36from lsst.pipe.base.testUtils import runTestQuantum, makeQuantum, assertValidOutput 

37 

38 

39class VisitConnections(PipelineTaskConnections, dimensions={"instrument", "visit"}): 

40 a = connectionTypes.Input( 

41 name="VisitA", 

42 storageClass="StructuredData", 

43 multiple=False, 

44 dimensions={"instrument", "visit"}, 

45 ) 

46 b = connectionTypes.Input( 

47 name="VisitB", 

48 storageClass="StructuredData", 

49 multiple=False, 

50 dimensions={"instrument", "visit"}, 

51 ) 

52 outA = connectionTypes.Output( 

53 name="VisitOutA", 

54 storageClass="StructuredData", 

55 multiple=False, 

56 dimensions={"instrument", "visit"}, 

57 ) 

58 outB = connectionTypes.Output( 

59 name="VisitOutB", 

60 storageClass="StructuredData", 

61 multiple=False, 

62 dimensions={"instrument", "visit"}, 

63 ) 

64 

65 

66class PatchConnections(PipelineTaskConnections, dimensions={"skymap", "tract"}): 

67 a = connectionTypes.Input( 

68 name="PatchA", 

69 storageClass="StructuredData", 

70 multiple=True, 

71 dimensions={"skymap", "tract", "patch"}, 

72 ) 

73 b = connectionTypes.Input( 

74 name="PatchB", 

75 storageClass="StructuredData", 

76 multiple=False, 

77 dimensions={"skymap", "tract"}, 

78 ) 

79 out = connectionTypes.Output( 

80 name="PatchOut", 

81 storageClass="StructuredData", 

82 multiple=True, 

83 dimensions={"skymap", "tract", "patch"}, 

84 ) 

85 

86 def __init__(self, *, config=None): 

87 super().__init__(config=config) 

88 

89 if not config.doUseB: 

90 self.inputs.remove("b") 

91 

92 

93class VisitConfig(PipelineTaskConfig, pipelineConnections=VisitConnections): 

94 pass 

95 

96 

97class PatchConfig(PipelineTaskConfig, pipelineConnections=PatchConnections): 

98 doUseB = lsst.pex.config.Field(default=True, dtype=bool, doc="") 

99 

100 

101class VisitTask(PipelineTask): 

102 ConfigClass = VisitConfig 

103 _DefaultName = "visit" 

104 

105 def run(self, a, b): 

106 outA = butlerTests.MetricsExample(data=(a.data + b.data)) 

107 outB = butlerTests.MetricsExample(data=(a.data * max(b.data))) 

108 return Struct(outA=outA, outB=outB) 

109 

110 

111class PatchTask(PipelineTask): 

112 ConfigClass = PatchConfig 

113 _DefaultName = "patch" 

114 

115 def run(self, a, b=None): 

116 if self.config.doUseB: 

117 out = [butlerTests.MetricsExample(data=(oneA.data + b.data)) for oneA in a] 

118 else: 

119 out = a 

120 return Struct(out=out) 

121 

122 

123class PipelineTaskTestSuite(lsst.utils.tests.TestCase): 

124 @classmethod 

125 def setUpClass(cls): 

126 super().setUpClass() 

127 # Repository should be re-created for each test case, but 

128 # this has a prohibitive run-time cost at present 

129 cls.root = tempfile.mkdtemp() 

130 

131 dataIds = { 

132 "instrument": ["notACam"], 

133 "physical_filter": ["k2020"], # needed for expandUniqueId(visit) 

134 "visit": [101, 102], 

135 "skymap": ["sky"], 

136 "tract": [42], 

137 "patch": [0, 1], 

138 } 

139 cls.repo = butlerTests.makeTestRepo(cls.root, dataIds) 

140 butlerTests.registerMetricsExample(cls.repo) 

141 

142 for typeName in {"VisitA", "VisitB", "VisitOutA", "VisitOutB"}: 

143 butlerTests.addDatasetType(cls.repo, typeName, {"instrument", "visit"}, "StructuredData") 

144 for typeName in {"PatchA", "PatchOut"}: 

145 butlerTests.addDatasetType(cls.repo, typeName, {"skymap", "tract", "patch"}, "StructuredData") 

146 butlerTests.addDatasetType(cls.repo, "PatchB", {"skymap", "tract"}, "StructuredData") 

147 

148 @classmethod 

149 def tearDownClass(cls): 

150 shutil.rmtree(cls.root, ignore_errors=True) 

151 super().tearDownClass() 

152 

153 def setUp(self): 

154 super().setUp() 

155 self.butler = butlerTests.makeTestCollection(self.repo) 

156 

157 def _makeVisitTestData(self, dataId): 

158 """Create dummy datasets suitable for VisitTask. 

159 

160 This method updates ``self.butler`` directly. 

161 

162 Parameters 

163 ---------- 

164 dataId : any data ID type 

165 The (shared) ID for the datasets to create. 

166 

167 Returns 

168 ------- 

169 datasets : `dict` [`str`, `list`] 

170 A dictionary keyed by dataset type. Its values are the list of 

171 integers used to create each dataset. The datasets stored in the 

172 butler are `lsst.daf.butler.tests.MetricsExample` objects with 

173 these lists as their ``data`` argument, but the lists are easier 

174 to manipulate in test code. 

175 """ 

176 inA = [1, 2, 3] 

177 inB = [4, 0, 1] 

178 self.butler.put(butlerTests.MetricsExample(data=inA), "VisitA", dataId) 

179 self.butler.put(butlerTests.MetricsExample(data=inB), "VisitB", dataId) 

180 return {"VisitA": inA, "VisitB": inB, } 

181 

182 def _makePatchTestData(self, dataId): 

183 """Create dummy datasets suitable for PatchTask. 

184 

185 This method updates ``self.butler`` directly. 

186 

187 Parameters 

188 ---------- 

189 dataId : any data ID type 

190 The (shared) ID for the datasets to create. Any patch ID is 

191 overridden to create multiple datasets. 

192 

193 Returns 

194 ------- 

195 datasets : `dict` [`str`, `list` [`tuple` [data ID, `list`]]] 

196 A dictionary keyed by dataset type. Its values are the data ID 

197 of each dataset and the list of integers used to create each. The 

198 datasets stored in the butler are 

199 `lsst.daf.butler.tests.MetricsExample` objects with these lists as 

200 their ``data`` argument, but the lists are easier to manipulate 

201 in test code. 

202 """ 

203 inA = [1, 2, 3] 

204 inB = [4, 0, 1] 

205 datasets = {"PatchA": [], "PatchB": []} 

206 for patch in {0, 1}: 

207 self.butler.put(butlerTests.MetricsExample(data=(inA + [patch])), "PatchA", dataId, patch=patch) 

208 datasets["PatchA"].append((dict(dataId, patch=patch), inA + [patch])) 

209 self.butler.put(butlerTests.MetricsExample(data=inB), "PatchB", dataId) 

210 datasets["PatchB"].append((dataId, inB)) 

211 return datasets 

212 

213 def testMakeQuantumNoSuchDatatype(self): 

214 config = VisitConfig() 

215 config.connections.a = "Visit" 

216 task = VisitTask(config=config) 

217 

218 dataId = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

219 self._makeVisitTestData(dataId) 

220 

221 with self.assertRaises(ValueError): 

222 makeQuantum(task, self.butler, {key: dataId for key in {"a", "b", "outA", "outB"}}) 

223 

224 def testMakeQuantumInvalidDimension(self): 

225 config = VisitConfig() 

226 config.connections.a = "PatchA" 

227 task = VisitTask(config=config) 

228 dataIdV = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

229 dataIdP = butlerTests.expandUniqueId(self.butler, {"patch": 0}) 

230 

231 inA = [1, 2, 3] 

232 inB = [4, 0, 1] 

233 self.butler.put(butlerTests.MetricsExample(data=inA), "VisitA", dataIdV) 

234 self.butler.put(butlerTests.MetricsExample(data=inA), "PatchA", dataIdP) 

235 self.butler.put(butlerTests.MetricsExample(data=inB), "VisitB", dataIdV) 

236 

237 with self.assertRaises(ValueError): 

238 makeQuantum(task, self.butler, { 

239 "a": dataIdP, 

240 "b": dataIdV, 

241 "outA": dataIdV, 

242 "outB": dataIdV, 

243 }) 

244 

245 def testMakeQuantumMissingMultiple(self): 

246 task = PatchTask() 

247 

248 dataId = butlerTests.expandUniqueId(self.butler, {"tract": 42}) 

249 self._makePatchTestData(dataId) 

250 

251 with self.assertRaises(ValueError): 

252 makeQuantum(task, self.butler, { 

253 "a": dict(dataId, patch=0), 

254 "b": dataId, 

255 "out": [dict(dataId, patch=patch) for patch in {0, 1}], 

256 }) 

257 

258 def testMakeQuantumExtraMultiple(self): 

259 task = PatchTask() 

260 

261 dataId = butlerTests.expandUniqueId(self.butler, {"tract": 42}) 

262 self._makePatchTestData(dataId) 

263 

264 with self.assertRaises(ValueError): 

265 makeQuantum(task, self.butler, { 

266 "a": [dict(dataId, patch=patch) for patch in {0, 1}], 

267 "b": [dataId], 

268 "out": [dict(dataId, patch=patch) for patch in {0, 1}], 

269 }) 

270 

271 def testMakeQuantumMissingDataId(self): 

272 task = VisitTask() 

273 

274 dataId = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

275 self._makeVisitTestData(dataId) 

276 

277 with self.assertRaises(ValueError): 

278 makeQuantum(task, self.butler, {key: dataId for key in {"a", "outA", "outB"}}) 

279 with self.assertRaises(ValueError): 

280 makeQuantum(task, self.butler, {key: dataId for key in {"a", "b", "outB"}}) 

281 

282 def testMakeQuantumCorruptedDataId(self): 

283 task = VisitTask() 

284 

285 dataId = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

286 self._makeVisitTestData(dataId) 

287 

288 with self.assertRaises(ValueError): 

289 # third argument should be a mapping keyed by component name 

290 makeQuantum(task, self.butler, dataId) 

291 

292 def testRunTestQuantumVisitWithRun(self): 

293 task = VisitTask() 

294 

295 dataId = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

296 data = self._makeVisitTestData(dataId) 

297 

298 quantum = makeQuantum(task, self.butler, {key: dataId for key in {"a", "b", "outA", "outB"}}) 

299 runTestQuantum(task, self.butler, quantum, mockRun=False) 

300 

301 # Can we use runTestQuantum to verify that task.run got called with correct inputs/outputs? 

302 self.assertTrue(self.butler.datasetExists("VisitOutA", dataId)) 

303 self.assertEqual(self.butler.get("VisitOutA", dataId), 

304 butlerTests.MetricsExample(data=(data["VisitA"] + data["VisitB"]))) 

305 self.assertTrue(self.butler.datasetExists("VisitOutB", dataId)) 

306 self.assertEqual(self.butler.get("VisitOutB", dataId), 

307 butlerTests.MetricsExample(data=(data["VisitA"] * max(data["VisitB"])))) 

308 

309 def testRunTestQuantumPatchWithRun(self): 

310 task = PatchTask() 

311 

312 dataId = butlerTests.expandUniqueId(self.butler, {"tract": 42}) 

313 data = self._makePatchTestData(dataId) 

314 

315 quantum = makeQuantum(task, self.butler, { 

316 "a": [dataset[0] for dataset in data["PatchA"]], 

317 "b": dataId, 

318 "out": [dataset[0] for dataset in data["PatchA"]], 

319 }) 

320 runTestQuantum(task, self.butler, quantum, mockRun=False) 

321 

322 # Can we use runTestQuantum to verify that task.run got called with correct inputs/outputs? 

323 inB = data["PatchB"][0][1] 

324 for dataset in data["PatchA"]: 

325 patchId = dataset[0] 

326 self.assertTrue(self.butler.datasetExists("PatchOut", patchId)) 

327 self.assertEqual(self.butler.get("PatchOut", patchId), 

328 butlerTests.MetricsExample(data=(dataset[1] + inB))) 

329 

330 def testRunTestQuantumVisitMockRun(self): 

331 task = VisitTask() 

332 

333 dataId = butlerTests.expandUniqueId(self.butler, {"visit": 102}) 

334 data = self._makeVisitTestData(dataId) 

335 

336 quantum = makeQuantum(task, self.butler, {key: dataId for key in {"a", "b", "outA", "outB"}}) 

337 run = runTestQuantum(task, self.butler, quantum, mockRun=True) 

338 

339 # Can we use the mock to verify that task.run got called with the correct inputs? 

340 run.assert_called_once_with(a=butlerTests.MetricsExample(data=data["VisitA"]), 

341 b=butlerTests.MetricsExample(data=data["VisitB"])) 

342 

343 def testRunTestQuantumPatchMockRun(self): 

344 task = PatchTask() 

345 

346 dataId = butlerTests.expandUniqueId(self.butler, {"tract": 42}) 

347 data = self._makePatchTestData(dataId) 

348 

349 quantum = makeQuantum(task, self.butler, { 

350 # Use lists, not sets, to ensure order agrees with test assertion 

351 "a": [dataset[0] for dataset in data["PatchA"]], 

352 "b": dataId, 

353 "out": [dataset[0] for dataset in data["PatchA"]], 

354 }) 

355 run = runTestQuantum(task, self.butler, quantum, mockRun=True) 

356 

357 # Can we use the mock to verify that task.run got called with the correct inputs? 

358 run.assert_called_once_with( 

359 a=[butlerTests.MetricsExample(data=dataset[1]) for dataset in data["PatchA"]], 

360 b=butlerTests.MetricsExample(data=data["PatchB"][0][1]) 

361 ) 

362 

363 def testRunTestQuantumPatchOptionalInput(self): 

364 config = PatchConfig() 

365 config.doUseB = False 

366 task = PatchTask(config=config) 

367 

368 dataId = butlerTests.expandUniqueId(self.butler, {"tract": 42}) 

369 data = self._makePatchTestData(dataId) 

370 

371 quantum = makeQuantum(task, self.butler, { 

372 # Use lists, not sets, to ensure order agrees with test assertion 

373 "a": [dataset[0] for dataset in data["PatchA"]], 

374 "out": [dataset[0] for dataset in data["PatchA"]], 

375 }) 

376 run = runTestQuantum(task, self.butler, quantum, mockRun=True) 

377 

378 # Can we use the mock to verify that task.run got called with the correct inputs? 

379 run.assert_called_once_with( 

380 a=[butlerTests.MetricsExample(data=dataset[1]) for dataset in data["PatchA"]] 

381 ) 

382 

383 def testAssertValidOutputPass(self): 

384 task = VisitTask() 

385 

386 inA = butlerTests.MetricsExample(data=[1, 2, 3]) 

387 inB = butlerTests.MetricsExample(data=[4, 0, 1]) 

388 result = task.run(inA, inB) 

389 

390 # should not throw 

391 assertValidOutput(task, result) 

392 

393 def testAssertValidOutputMissing(self): 

394 task = VisitTask() 

395 

396 def run(a, b): 

397 return Struct(outA=a) 

398 task.run = run 

399 

400 inA = butlerTests.MetricsExample(data=[1, 2, 3]) 

401 inB = butlerTests.MetricsExample(data=[4, 0, 1]) 

402 result = task.run(inA, inB) 

403 

404 with self.assertRaises(AssertionError): 

405 assertValidOutput(task, result) 

406 

407 def testAssertValidOutputSingle(self): 

408 task = PatchTask() 

409 

410 def run(a, b): 

411 return Struct(out=b) 

412 task.run = run 

413 

414 inA = butlerTests.MetricsExample(data=[1, 2, 3]) 

415 inB = butlerTests.MetricsExample(data=[4, 0, 1]) 

416 result = task.run([inA], inB) 

417 

418 with self.assertRaises(AssertionError): 

419 assertValidOutput(task, result) 

420 

421 def testAssertValidOutputMultiple(self): 

422 task = VisitTask() 

423 

424 def run(a, b): 

425 return Struct(outA=[a], outB=b) 

426 task.run = run 

427 

428 inA = butlerTests.MetricsExample(data=[1, 2, 3]) 

429 inB = butlerTests.MetricsExample(data=[4, 0, 1]) 

430 result = task.run(inA, inB) 

431 

432 with self.assertRaises(AssertionError): 

433 assertValidOutput(task, result) 

434 

435 

436class MyMemoryTestCase(lsst.utils.tests.MemoryTestCase): 

437 pass 

438 

439 

440def setup_module(module): 

441 lsst.utils.tests.init() 

442 

443 

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

445 lsst.utils.tests.init() 

446 unittest.main()