Coverage for tests/test_cliCmdQueryCollections.py: 24%

112 statements  

« prev     ^ index     » next       coverage.py v7.11.0, created at 2025-11-19 01:32 -0800

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.utils import ButlerTestHelper, readTable 

42 

43 

44class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

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

46 

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

48 

49 @staticmethod 

50 def defaultExpected(): 

51 return dict( 

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

53 ) 

54 

55 @staticmethod 

56 def command(): 

57 return query_collections 

58 

59 def test_minimal(self): 

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

61 self.run_test( 

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

63 self.makeExpected( 

64 repo="here", 

65 show_dataset_types=False, 

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

67 ), 

68 ) 

69 

70 def test_all(self): 

71 """Test all parameters""" 

72 self.run_test( 

73 [ 

74 "query-collections", 

75 "here", 

76 "foo*", 

77 "--collection-type", 

78 "TAGGED", 

79 "--collection-type", 

80 "RUN", 

81 "--chains", 

82 "TABLE", 

83 ], 

84 self.makeExpected( 

85 repo="here", 

86 glob=("foo*",), 

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

88 chains="TABLE", 

89 show_dataset_types=False, 

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

91 ), 

92 ) 

93 

94 def test_show_dataset_types(self): 

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

96 self.run_test( 

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

98 self.makeExpected( 

99 repo="here", 

100 chains="TABLE", 

101 show_dataset_types=True, 

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

103 ), 

104 ) 

105 self.run_test( 

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

107 self.makeExpected( 

108 repo="here", 

109 chains="TREE", 

110 show_dataset_types=True, 

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

112 ), 

113 ) 

114 self.run_test( 

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

116 self.makeExpected( 

117 repo="here", 

118 chains="FLATTEN", 

119 show_dataset_types=True, 

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

121 ), 

122 ) 

123 

124 def test_exclude_dataset_types(self): 

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

126 self.run_test( 

127 [ 

128 "query-collections", 

129 "here", 

130 "--chains", 

131 "TABLE", 

132 "--show-dataset-types", 

133 "--exclude-dataset-types", 

134 "flat", 

135 ], 

136 self.makeExpected( 

137 repo="here", 

138 show_dataset_types=True, 

139 exclude_dataset_types=("flat",), 

140 ), 

141 ) 

142 

143 

144class QueryCollectionsScriptTest(ButlerTestHelper, unittest.TestCase): 

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

146 

147 def setUp(self): 

148 self.runner = LogCliRunner() 

149 

150 def testGetCollections(self): 

151 run = "ingest/run" 

152 tag = "tag" 

153 with self.runner.isolated_filesystem(): 

154 butlerCfg = Butler.makeRepo("here") 

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

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

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

158 

159 # Verify collections that were created are found by 

160 # query-collections. 

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

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

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

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

165 

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

167 # name matches with the specified pattern are returned. 

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

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

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

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

172 

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

174 # that type are returned. 

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

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

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

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

179 

180 

181class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

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

183 

184 def setUp(self): 

185 self.runner = LogCliRunner() 

186 

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

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

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

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

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

192 

193 def testChained(self): 

194 with self.runner.isolated_filesystem(): 

195 # Create a butler and add some chained collections: 

196 butlerCfg = Butler.makeRepo("here") 

197 

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

199 

200 # Replace datastore functions with mocks: 

201 DatastoreMock.apply(butler1) 

202 

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

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

205 registry1 = butler1.registry 

206 registry1.registerRun("run1") 

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

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

209 

210 # Create the collection chain 

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

212 self.assertChain( 

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

214 ) 

215 

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

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

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

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

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

221 # the command line interface. 

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

223 

224 expected = Table( 

225 array( 

226 ( 

227 ("calibration1", "CALIBRATION"), 

228 ("chain1", "CHAINED"), 

229 (" tag1", "TAGGED"), 

230 (" run1", "RUN"), 

231 (" chain2", "CHAINED"), 

232 (" calibration1", "CALIBRATION"), 

233 (" run1", "RUN"), 

234 ("chain2", "CHAINED"), 

235 (" calibration1", "CALIBRATION"), 

236 (" run1", "RUN"), 

237 ("imported_g", "RUN"), 

238 ("imported_r", "RUN"), 

239 ("run1", "RUN"), 

240 ("tag1", "TAGGED"), 

241 ) 

242 ), 

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

244 ) 

245 self.assertAstropyTablesEqual(table, expected) 

246 

247 # Test table with inverse == True 

248 table = queryCollections( 

249 "here", 

250 glob=(), 

251 collection_type=CollectionType.all(), 

252 chains="INVERSE-TREE", 

253 ) 

254 expected = Table( 

255 array( 

256 ( 

257 ("calibration1", "CALIBRATION"), 

258 (" chain2", "CHAINED"), 

259 (" chain1", "CHAINED"), 

260 ("chain1", "CHAINED"), 

261 ("chain2", "CHAINED"), 

262 (" chain1", "CHAINED"), 

263 ("imported_g", "RUN"), 

264 ("imported_r", "RUN"), 

265 ("run1", "RUN"), 

266 (" chain1", "CHAINED"), 

267 (" chain2", "CHAINED"), 

268 (" chain1", "CHAINED"), 

269 ("tag1", "TAGGED"), 

270 (" chain1", "CHAINED"), 

271 ) 

272 ), 

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

274 ) 

275 self.assertAstropyTablesEqual(table, expected) 

276 

277 # Test table with show_dataset_types == True 

278 table = queryCollections( 

279 "here", 

280 glob=(), 

281 collection_type=CollectionType.all(), 

282 chains="TREE", 

283 show_dataset_types=True, 

284 ) 

285 expected = Table( 

286 array( 

287 ( 

288 ("calibration1", "CALIBRATION", ""), 

289 ("chain1", "CHAINED", ""), 

290 (" tag1", "TAGGED", ""), 

291 (" run1", "RUN", ""), 

292 (" chain2", "CHAINED", ""), 

293 (" calibration1", "CALIBRATION", ""), 

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

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

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

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

298 ("imported_g", "RUN", "bias"), 

299 ("", "", "flat"), 

300 ("imported_r", "RUN", "bias"), 

301 ("", "", "flat"), 

302 ("run1", "RUN", ""), 

303 ("tag1", "TAGGED", ""), 

304 ) 

305 ), 

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

307 ) 

308 self.assertAstropyTablesEqual(table, expected) 

309 

310 # Test table with show_dataset_types, excluding some dataset types 

311 table = queryCollections( 

312 "here", 

313 glob=(), 

314 collection_type=CollectionType.all(), 

315 chains="TREE", 

316 show_dataset_types=True, 

317 exclude_dataset_types=("bias",), 

318 ) 

319 expected = Table( 

320 array( 

321 ( 

322 ("calibration1", "CALIBRATION", ""), 

323 ("chain1", "CHAINED", ""), 

324 (" tag1", "TAGGED", ""), 

325 (" run1", "RUN", ""), 

326 (" chain2", "CHAINED", ""), 

327 (" calibration1", "CALIBRATION", ""), 

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

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

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

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

332 ("imported_g", "RUN", "flat"), 

333 ("imported_r", "RUN", "flat"), 

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

335 ("tag1", "TAGGED", ""), 

336 ) 

337 ), 

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

339 ) 

340 self.assertAstropyTablesEqual(table, expected) 

341 

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

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

344 expected = Table( 

345 array( 

346 ( 

347 ("calibration1", "CALIBRATION", ""), 

348 ("chain1", "CHAINED", "tag1"), 

349 ("", "", "run1"), 

350 ("", "", "chain2"), 

351 ("chain2", "CHAINED", "calibration1"), 

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

353 ("imported_g", "RUN", ""), 

354 ("imported_r", "RUN", ""), 

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

356 ("tag1", "TAGGED", ""), 

357 ) 

358 ), 

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

360 ) 

361 table = readTable(result.output) 

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

363 

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

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

366 expected = Table( 

367 array( 

368 ( 

369 ("calibration1", "CALIBRATION", "chain2"), 

370 ("chain1", "CHAINED", ""), 

371 ("chain2", "CHAINED", "chain1"), 

372 ("imported_g", "RUN", ""), 

373 ("imported_r", "RUN", ""), 

374 ("run1", "RUN", "chain1"), 

375 ("", "", "chain2"), 

376 ("tag1", "TAGGED", "chain1"), 

377 ) 

378 ), 

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

380 ) 

381 table = readTable(result.output) 

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

383 

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

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

386 expected = Table( 

387 array( 

388 ( 

389 ("calibration1", "CALIBRATION"), 

390 ("imported_g", "RUN"), 

391 ("imported_r", "RUN"), 

392 ("run1", "RUN"), 

393 ("tag1", "TAGGED"), 

394 ) 

395 ), 

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

397 ) 

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

399 

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

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

402 expected = Table( 

403 array( 

404 ( 

405 ("calibration1", "CALIBRATION"), 

406 ("chain1", "CHAINED"), 

407 ("chain2", "CHAINED"), 

408 ("imported_g", "RUN"), 

409 ("imported_r", "RUN"), 

410 ("run1", "RUN"), 

411 ("tag1", "TAGGED"), 

412 ) 

413 ), 

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

415 ) 

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

417 

418 # Add a couple more run collections for chain testing 

419 registry1.registerRun("run2") 

420 registry1.registerRun("run3") 

421 registry1.registerRun("run4") 

422 

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

424 

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

426 

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

428 

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

430 

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

432 

433 self.assertChain( 

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

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

436 ) 

437 

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

439 

440 # Out-of-bounds index 

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

442 self.assertEqual(result.exit_code, 1) 

443 

444 

445if __name__ == "__main__": 

446 unittest.main()