Coverage for tests/test_cliCmdQueryCollections.py: 100%

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 software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

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

29 

30import unittest 

31 

32from astropy.table import Table 

33from numpy import array 

34 

35from lsst.daf.butler import Butler, CollectionType 

36from lsst.daf.butler.cli.butler import cli 

37from lsst.daf.butler.cli.cmd import query_collections 

38from lsst.daf.butler.cli.utils import LogCliRunner, clickResultMsg 

39from lsst.daf.butler.script import queryCollections 

40from lsst.daf.butler.tests import CliCmdTestBase, DatastoreMock 

41from lsst.daf.butler.tests._repo_template_cache import make_repo_for_test 

42from lsst.daf.butler.tests.utils import ButlerTestHelper, readTable 

43 

44 

45class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

46 """Test the query-collections command-line.""" 

47 

48 mockFuncName = "lsst.daf.butler.cli.cmd.commands.script.queryCollections" 

49 

50 @staticmethod 

51 def defaultExpected(): 

52 return dict( 

53 repo=None, collection_type=tuple(CollectionType.__members__.values()), chains="TABLE", glob=() 

54 ) 

55 

56 @staticmethod 

57 def command(): 

58 return query_collections 

59 

60 def test_minimal(self): 

61 """Test only required parameters, and omit optional parameters.""" 

62 self.run_test( 

63 ["query-collections", "here", "--chains", "TABLE"], 

64 self.makeExpected( 

65 repo="here", 

66 show_dataset_types=False, 

67 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"), 

68 ), 

69 ) 

70 

71 def test_all(self): 

72 """Test all parameters""" 

73 self.run_test( 

74 [ 

75 "query-collections", 

76 "here", 

77 "foo*", 

78 "--collection-type", 

79 "TAGGED", 

80 "--collection-type", 

81 "RUN", 

82 "--chains", 

83 "TABLE", 

84 ], 

85 self.makeExpected( 

86 repo="here", 

87 glob=("foo*",), 

88 collection_type=(CollectionType.TAGGED, CollectionType.RUN), 

89 chains="TABLE", 

90 show_dataset_types=False, 

91 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"), 

92 ), 

93 ) 

94 

95 def test_show_dataset_types(self): 

96 """Test the --show-dataset-types option.""" 

97 self.run_test( 

98 ["query-collections", "here", "--chains", "TABLE", "--show-dataset-types"], 

99 self.makeExpected( 

100 repo="here", 

101 chains="TABLE", 

102 show_dataset_types=True, 

103 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"), 

104 ), 

105 ) 

106 self.run_test( 

107 ["query-collections", "here", "--chains", "TREE", "--show-dataset-types"], 

108 self.makeExpected( 

109 repo="here", 

110 chains="TREE", 

111 show_dataset_types=True, 

112 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"), 

113 ), 

114 ) 

115 self.run_test( 

116 ["query-collections", "here", "--chains", "FLATTEN", "--show-dataset-types"], 

117 self.makeExpected( 

118 repo="here", 

119 chains="FLATTEN", 

120 show_dataset_types=True, 

121 exclude_dataset_types=("*_config", "*_log", "*_metadata", "packages"), 

122 ), 

123 ) 

124 

125 def test_exclude_dataset_types(self): 

126 """Test the --exclude-dataset-types option.""" 

127 self.run_test( 

128 [ 

129 "query-collections", 

130 "here", 

131 "--chains", 

132 "TABLE", 

133 "--show-dataset-types", 

134 "--exclude-dataset-types", 

135 "flat", 

136 ], 

137 self.makeExpected( 

138 repo="here", 

139 show_dataset_types=True, 

140 exclude_dataset_types=("flat",), 

141 ), 

142 ) 

143 

144 

145class QueryCollectionsScriptTest(ButlerTestHelper, unittest.TestCase): 

146 """Test the query-collections script interface.""" 

147 

148 def setUp(self): 

149 self.runner = LogCliRunner() 

150 

151 def testGetCollections(self): 

152 run = "ingest/run" 

153 tag = "tag" 

154 with self.runner.isolated_filesystem(): 

155 butlerCfg = make_repo_for_test("here") 

156 # the purpose of this call is to create some collections 

157 butler = Butler.from_config(butlerCfg, run=run, collections=[tag], writeable=True) 

158 self.enterContext(butler) 

159 butler.registry.registerCollection(tag, CollectionType.TAGGED) 

160 

161 # Verify collections that were created are found by 

162 # query-collections. 

163 result = self.runner.invoke(cli, ["query-collections", "here"]) 

164 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

165 expected = Table((("ingest/run", "tag"), ("RUN", "TAGGED")), names=("Name", "Type")) 

166 self.assertAstropyTablesEqual(readTable(result.output), expected) 

167 

168 # Verify that with a glob argument, that only collections whose 

169 # name matches with the specified pattern are returned. 

170 result = self.runner.invoke(cli, ["query-collections", "here", "t*"]) 

171 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

172 expected = Table((("tag",), ("TAGGED",)), names=("Name", "Type")) 

173 self.assertAstropyTablesEqual(readTable(result.output), expected) 

174 

175 # Verify that with a collection type argument, only collections of 

176 # that type are returned. 

177 result = self.runner.invoke(cli, ["query-collections", "here", "--collection-type", "RUN"]) 

178 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

179 expected = Table((("ingest/run",), ("RUN",)), names=("Name", "Type")) 

180 self.assertAstropyTablesEqual(readTable(result.output), expected) 

181 

182 

183class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

184 """Test the collection-chain command-line interface.""" 

185 

186 def setUp(self): 

187 self.runner = LogCliRunner() 

188 

189 def assertChain(self, args: list[str], expected: str): 

190 """Run collection-chain and check the expected result""" 

191 result = self.runner.invoke(cli, ["collection-chain", "here", *args]) 

192 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

193 self.assertEqual(result.output.strip(), expected, clickResultMsg(result)) 

194 

195 def testChained(self): 

196 with self.runner.isolated_filesystem(): 

197 # Create a butler and add some chained collections: 

198 butlerCfg = make_repo_for_test("here") 

199 

200 butler1 = Butler.from_config(butlerCfg, writeable=True) 

201 self.enterContext(butler1) 

202 

203 # Replace datastore functions with mocks: 

204 DatastoreMock.apply(butler1) 

205 

206 butler1.import_(filename="resource://lsst.daf.butler/tests/registry_data/base.yaml") 

207 butler1.import_(filename="resource://lsst.daf.butler/tests/registry_data/datasets.yaml") 

208 registry1 = butler1.registry 

209 registry1.registerRun("run1") 

210 registry1.registerCollection("tag1", CollectionType.TAGGED) 

211 registry1.registerCollection("calibration1", CollectionType.CALIBRATION) 

212 

213 # Create the collection chain 

214 self.assertChain(["chain2", "calibration1", "run1"], "[calibration1, run1]") 

215 self.assertChain( 

216 ["--mode", "redefine", "chain1", "tag1", "run1", "chain2"], "[tag1, run1, chain2]" 

217 ) 

218 

219 # Use the script function to test the query-collections TREE 

220 # option, because the astropy.table.Table.read method, which we are 

221 # using for verification elsewhere in this file, seems to strip 

222 # leading whitespace from columns. This makes it impossible to test 

223 # the nested TREE output of the query-collections subcommand from 

224 # the command line interface. 

225 table = queryCollections("here", glob=(), collection_type=CollectionType.all(), chains="TREE") 

226 

227 expected = Table( 

228 array( 

229 ( 

230 ("calibration1", "CALIBRATION"), 

231 ("chain1", "CHAINED"), 

232 (" tag1", "TAGGED"), 

233 (" run1", "RUN"), 

234 (" chain2", "CHAINED"), 

235 (" calibration1", "CALIBRATION"), 

236 (" run1", "RUN"), 

237 ("chain2", "CHAINED"), 

238 (" calibration1", "CALIBRATION"), 

239 (" run1", "RUN"), 

240 ("imported_g", "RUN"), 

241 ("imported_r", "RUN"), 

242 ("run1", "RUN"), 

243 ("tag1", "TAGGED"), 

244 ) 

245 ), 

246 names=("Name", "Type"), 

247 ) 

248 self.assertAstropyTablesEqual(table, expected) 

249 

250 # Test table with inverse == True 

251 table = queryCollections( 

252 "here", 

253 glob=(), 

254 collection_type=CollectionType.all(), 

255 chains="INVERSE-TREE", 

256 ) 

257 expected = Table( 

258 array( 

259 ( 

260 ("calibration1", "CALIBRATION"), 

261 (" chain2", "CHAINED"), 

262 (" chain1", "CHAINED"), 

263 ("chain1", "CHAINED"), 

264 ("chain2", "CHAINED"), 

265 (" chain1", "CHAINED"), 

266 ("imported_g", "RUN"), 

267 ("imported_r", "RUN"), 

268 ("run1", "RUN"), 

269 (" chain1", "CHAINED"), 

270 (" chain2", "CHAINED"), 

271 (" chain1", "CHAINED"), 

272 ("tag1", "TAGGED"), 

273 (" chain1", "CHAINED"), 

274 ) 

275 ), 

276 names=("Name", "Type"), 

277 ) 

278 self.assertAstropyTablesEqual(table, expected) 

279 

280 # Test table with show_dataset_types == True 

281 table = queryCollections( 

282 "here", 

283 glob=(), 

284 collection_type=CollectionType.all(), 

285 chains="TREE", 

286 show_dataset_types=True, 

287 ) 

288 expected = Table( 

289 array( 

290 ( 

291 ("calibration1", "CALIBRATION", ""), 

292 ("chain1", "CHAINED", ""), 

293 (" tag1", "TAGGED", ""), 

294 (" run1", "RUN", ""), 

295 (" chain2", "CHAINED", ""), 

296 (" calibration1", "CALIBRATION", ""), 

297 (" run1", "RUN", ""), 

298 ("chain2", "CHAINED", ""), 

299 (" calibration1", "CALIBRATION", ""), 

300 (" run1", "RUN", ""), 

301 ("imported_g", "RUN", "bias"), 

302 ("", "", "flat"), 

303 ("imported_r", "RUN", "bias"), 

304 ("", "", "flat"), 

305 ("run1", "RUN", ""), 

306 ("tag1", "TAGGED", ""), 

307 ) 

308 ), 

309 names=("Name", "Type", "Dataset Types"), 

310 ) 

311 self.assertAstropyTablesEqual(table, expected) 

312 

313 # Test table with show_dataset_types, excluding some dataset types 

314 table = queryCollections( 

315 "here", 

316 glob=(), 

317 collection_type=CollectionType.all(), 

318 chains="TREE", 

319 show_dataset_types=True, 

320 exclude_dataset_types=("bias",), 

321 ) 

322 expected = Table( 

323 array( 

324 ( 

325 ("calibration1", "CALIBRATION", ""), 

326 ("chain1", "CHAINED", ""), 

327 (" tag1", "TAGGED", ""), 

328 (" run1", "RUN", ""), 

329 (" chain2", "CHAINED", ""), 

330 (" calibration1", "CALIBRATION", ""), 

331 (" run1", "RUN", ""), 

332 ("chain2", "CHAINED", ""), 

333 (" calibration1", "CALIBRATION", ""), 

334 (" run1", "RUN", ""), 

335 ("imported_g", "RUN", "flat"), 

336 ("imported_r", "RUN", "flat"), 

337 ("run1", "RUN", ""), 

338 ("tag1", "TAGGED", ""), 

339 ) 

340 ), 

341 names=("Name", "Type", "Dataset Types"), 

342 ) 

343 self.assertAstropyTablesEqual(table, expected) 

344 

345 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "TABLE"]) 

346 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

347 expected = Table( 

348 array( 

349 ( 

350 ("calibration1", "CALIBRATION", ""), 

351 ("chain1", "CHAINED", "tag1"), 

352 ("", "", "run1"), 

353 ("", "", "chain2"), 

354 ("chain2", "CHAINED", "calibration1"), 

355 ("", "", "run1"), 

356 ("imported_g", "RUN", ""), 

357 ("imported_r", "RUN", ""), 

358 ("run1", "RUN", ""), 

359 ("tag1", "TAGGED", ""), 

360 ) 

361 ), 

362 names=("Name", "Type", "Children"), 

363 ) 

364 table = readTable(result.output) 

365 self.assertAstropyTablesEqual(readTable(result.output), expected) 

366 

367 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "INVERSE-TABLE"]) 

368 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

369 expected = Table( 

370 array( 

371 ( 

372 ("calibration1", "CALIBRATION", "chain2"), 

373 ("chain1", "CHAINED", ""), 

374 ("chain2", "CHAINED", "chain1"), 

375 ("imported_g", "RUN", ""), 

376 ("imported_r", "RUN", ""), 

377 ("run1", "RUN", "chain1"), 

378 ("", "", "chain2"), 

379 ("tag1", "TAGGED", "chain1"), 

380 ) 

381 ), 

382 names=("Name", "Type", "Parents"), 

383 ) 

384 table = readTable(result.output) 

385 self.assertAstropyTablesEqual(readTable(result.output), expected) 

386 

387 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "FLATTEN"]) 

388 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

389 expected = Table( 

390 array( 

391 ( 

392 ("calibration1", "CALIBRATION"), 

393 ("imported_g", "RUN"), 

394 ("imported_r", "RUN"), 

395 ("run1", "RUN"), 

396 ("tag1", "TAGGED"), 

397 ) 

398 ), 

399 names=("Name", "Type"), 

400 ) 

401 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True) 

402 

403 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "NO-CHILDREN"]) 

404 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

405 expected = Table( 

406 array( 

407 ( 

408 ("calibration1", "CALIBRATION"), 

409 ("chain1", "CHAINED"), 

410 ("chain2", "CHAINED"), 

411 ("imported_g", "RUN"), 

412 ("imported_r", "RUN"), 

413 ("run1", "RUN"), 

414 ("tag1", "TAGGED"), 

415 ) 

416 ), 

417 names=("Name", "Type"), 

418 ) 

419 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True) 

420 

421 # Add a couple more run collections for chain testing 

422 registry1.registerRun("run2") 

423 registry1.registerRun("run3") 

424 registry1.registerRun("run4") 

425 

426 self.assertChain(["--mode", "pop", "chain1"], "[run1, chain2]") 

427 

428 self.assertChain(["--mode", "extend", "chain1", "run2", "run3"], "[run1, chain2, run2, run3]") 

429 

430 self.assertChain(["--mode", "remove", "chain1", "chain2", "run2"], "[run1, run3]") 

431 

432 self.assertChain(["--mode", "prepend", "chain1", "chain2", "run2"], "[chain2, run2, run1, run3]") 

433 

434 self.assertChain(["--mode", "pop", "chain1", "1", "3"], "[chain2, run1]") 

435 

436 self.assertChain( 

437 ["--mode", "redefine", "chain1", "chain2", "run2", "run3,run4", "--flatten"], 

438 "[calibration1, run1, run2, run3, run4]", 

439 ) 

440 

441 self.assertChain(["--mode", "pop", "chain1", "--", "-1", "-3"], "[calibration1, run1, run3]") 

442 

443 # Out-of-bounds index 

444 result = self.runner.invoke(cli, ["collection-chain", "here", "--mode", "pop", "chain1", "10"]) 

445 self.assertEqual(result.exit_code, 1) 

446 

447 

448if __name__ == "__main__": 

449 unittest.main()