Coverage for tests/test_cliCmdPruneDatasets.py: 47%

106 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-25 10:24 -0700

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 prune-datasets subcommand. 

29""" 

30 

31import unittest 

32from unittest.mock import patch 

33 

34# Tests require the SqlRegistry 

35import lsst.daf.butler.registry.sql_registry 

36import lsst.daf.butler.script 

37from astropy.table import Table 

38from lsst.daf.butler.cli.butler import cli as butlerCli 

39from lsst.daf.butler.cli.cmd.commands import ( 

40 pruneDatasets_askContinueMsg, 

41 pruneDatasets_didNotRemoveAforementioned, 

42 pruneDatasets_didRemoveAforementioned, 

43 pruneDatasets_didRemoveMsg, 

44 pruneDatasets_errNoCollectionRestriction, 

45 pruneDatasets_errNoOp, 

46 pruneDatasets_errPruneOnNotRun, 

47 pruneDatasets_errPurgeAndDisassociate, 

48 pruneDatasets_errQuietWithDryRun, 

49 pruneDatasets_noDatasetsFound, 

50 pruneDatasets_willRemoveMsg, 

51 pruneDatasets_wouldDisassociateAndRemoveMsg, 

52 pruneDatasets_wouldDisassociateMsg, 

53 pruneDatasets_wouldRemoveMsg, 

54) 

55from lsst.daf.butler.cli.utils import LogCliRunner, astropyTablesToStr, clickResultMsg 

56from lsst.daf.butler.direct_butler import DirectButler 

57from lsst.daf.butler.registry import CollectionType 

58from lsst.daf.butler.script import QueryDatasets 

59 

60doFindTables = True 

61 

62 

63def getTables(): 

64 """Return test table.""" 

65 if doFindTables: 

66 return (Table(((1, 2, 3),), names=("foo",)),) 

67 return () 

68 

69 

70def getDatasets(): 

71 """Return the datasets string.""" 

72 return "datasets" 

73 

74 

75def makeQueryDatasets(*args, **kwargs): 

76 """Return a query datasets object.""" 

77 return QueryDatasets(*args, **kwargs) 

78 

79 

80class PruneDatasetsTestCase(unittest.TestCase): 

81 """Tests the ``prune_datasets`` "command" function (in 

82 ``cli/cmd/commands.py``) and the ``pruneDatasets`` "script" function (in 

83 ``scripts/_pruneDatasets.py``). 

84 

85 ``Butler.pruneDatasets`` and a few other functions that get called before 

86 it are mocked, and tests check for expected arguments to those mocks. 

87 """ 

88 

89 def setUp(self): 

90 self.repo = "here" 

91 

92 @staticmethod 

93 def makeQueryDatasetsArgs(*, repo, **kwargs): 

94 expectedArgs = dict( 

95 repo=repo, collections=(), where="", find_first=True, show_uri=False, glob=tuple() 

96 ) 

97 expectedArgs.update(kwargs) 

98 return expectedArgs 

99 

100 @staticmethod 

101 def makePruneDatasetsArgs(**kwargs): 

102 expectedArgs = dict(refs=tuple(), disassociate=False, tags=(), purge=False, unstore=False) 

103 expectedArgs.update(kwargs) 

104 return expectedArgs 

105 

106 # Mock the QueryDatasets.getTables function to return a set of Astropy 

107 # tables, similar to what would be returned by a call to 

108 # QueryDatasets.getTables on a repo with real data. 

109 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getTables", side_effect=getTables) 

110 # Mock the QueryDatasets.getDatasets function. Normally it would return a 

111 # list of queries.DatasetQueryResults, but all we need to do is verify that 

112 # the output of this function is passed into our pruneDatasets magicMock, 

113 # so we can return something arbitrary that we can test is equal.""" 

114 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getDatasets", side_effect=getDatasets) 

115 # Mock the actual QueryDatasets class, so we can inspect calls to its init 

116 # function. Note that the side_effect returns an instance of QueryDatasets, 

117 # so this mock records and then is a pass-through. 

118 @patch.object(lsst.daf.butler.script._pruneDatasets, "QueryDatasets", side_effect=makeQueryDatasets) 

119 # Mock the pruneDatasets butler command so we can test for expected calls 

120 # to it, without dealing with setting up a full repo with data for it. 

121 @patch.object(DirectButler, "pruneDatasets") 

122 def run_test( 

123 self, 

124 mockPruneDatasets, 

125 mockQueryDatasets_init, 

126 mockQueryDatasets_getDatasets, 

127 mockQueryDatasets_getTables, 

128 cliArgs, 

129 exMsgs, 

130 exPruneDatasetsCallArgs, 

131 exGetTablesCalled, 

132 exQueryDatasetsCallArgs, 

133 invokeInput=None, 

134 exPruneDatasetsExitCode=0, 

135 ): 

136 """Execute the test. 

137 

138 Makes a temporary repo, invokes ``prune-datasets``. Verifies expected 

139 output, exit codes, and mock calls. 

140 

141 Parameters 

142 ---------- 

143 mockPruneDatasets : `MagicMock` 

144 The MagicMock for the ``Butler.pruneDatasets`` function. 

145 mockQueryDatasets_init : `MagicMock` 

146 The MagicMock for the ``QueryDatasets.__init__`` function. 

147 mockQueryDatasets_getDatasets : `MagicMock` 

148 The MagicMock for the ``QueryDatasets.getDatasets`` function. 

149 mockQueryDatasets_getTables : `MagicMock` 

150 The MagicMock for the ``QueryDatasets.getTables`` function. 

151 cliArgs : `list` [`str`] 

152 The arguments to pass to the command line. Do not include the 

153 subcommand name or the repo. 

154 exMsgs : `list` [`str`] or None 

155 A list of text fragments that should appear in the text output 

156 after calling the CLI command, or None if no output should be 

157 produced. 

158 exPruneDatasetsCallArgs : `dict` [`str`, `Any`] 

159 The arguments that ``Butler.pruneDatasets`` should have been called 

160 with, or None if that function should not have been called. 

161 exGetTablesCalled : bool 

162 `True` if ``QueryDatasets.getTables`` should have been called, else 

163 `False`. 

164 exQueryDatasetsCallArgs : `dict` [`str`, `Any`] 

165 The arguments that ``QueryDatasets.__init__`` should have bene 

166 called with, or `None` if the function should not have been called. 

167 invokeInput : `str`, optional. 

168 As string to pass to the ``CliRunner.invoke`` `input` argument. By 

169 default None. 

170 exPruneDatasetsExitCode : `int` 

171 The expected exit code returned from invoking ``prune-datasets``. 

172 """ 

173 runner = LogCliRunner() 

174 with runner.isolated_filesystem(): 

175 # Make a repo so a butler can be created 

176 result = runner.invoke(butlerCli, ["create", self.repo]) 

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

178 

179 # Run the prune-datasets CLI command, this will call all of our 

180 # mocks: 

181 cliArgs = ["prune-datasets", self.repo] + cliArgs 

182 result = runner.invoke(butlerCli, cliArgs, input=invokeInput) 

183 self.assertEqual(result.exit_code, exPruneDatasetsExitCode, clickResultMsg(result)) 

184 

185 # Verify the Butler.pruneDatasets was called exactly once with 

186 # expected arguments. The datasets argument is the value returned 

187 # by QueryDatasets, which we've mocked with side effect 

188 # ``getDatasets()``. 

189 if exPruneDatasetsCallArgs: 

190 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs) 

191 else: 

192 mockPruneDatasets.assert_not_called() 

193 

194 # Less critical, but do a quick verification that the QueryDataset 

195 # member function mocks were called, in this case we expect one 

196 # time each. 

197 if exQueryDatasetsCallArgs: 

198 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs) 

199 else: 

200 mockQueryDatasets_init.assert_not_called() 

201 # If Butler.pruneDatasets was not called, then 

202 # QueryDatasets.getDatasets also does not get called. 

203 if exPruneDatasetsCallArgs: 

204 mockQueryDatasets_getDatasets.assert_called_once() 

205 else: 

206 mockQueryDatasets_getDatasets.assert_not_called() 

207 if exGetTablesCalled: 

208 mockQueryDatasets_getTables.assert_called_once() 

209 else: 

210 mockQueryDatasets_getTables.assert_not_called() 

211 

212 if exMsgs is None: 

213 self.assertEqual("", result.output) 

214 else: 

215 for expectedMsg in exMsgs: 

216 self.assertIn(expectedMsg, result.output) 

217 

218 def test_defaults_doContinue(self): 

219 """Test running with the default values. 

220 

221 Verify that with the default flags that the subcommand says what it 

222 will do, prompts for input, and says that it's done. 

223 """ 

224 self.run_test( 

225 cliArgs=["myCollection", "--unstore"], 

226 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True), 

227 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)), 

228 exGetTablesCalled=True, 

229 exMsgs=( 

230 pruneDatasets_willRemoveMsg, 

231 pruneDatasets_askContinueMsg, 

232 astropyTablesToStr(getTables()), 

233 pruneDatasets_didRemoveAforementioned, 

234 ), 

235 invokeInput="yes", 

236 ) 

237 

238 def test_defaults_doNotContinue(self): 

239 """Test running with the default values but not continuing. 

240 

241 Verify that with the default flags that the subcommand says what it 

242 will do, prompts for input, and aborts when told not to continue. 

243 """ 

244 self.run_test( 

245 cliArgs=["myCollection", "--unstore"], 

246 exPruneDatasetsCallArgs=None, 

247 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)), 

248 exGetTablesCalled=True, 

249 exMsgs=( 

250 pruneDatasets_willRemoveMsg, 

251 pruneDatasets_askContinueMsg, 

252 pruneDatasets_didNotRemoveAforementioned, 

253 ), 

254 invokeInput="no", 

255 ) 

256 

257 def test_dryRun_unstore(self): 

258 """Test the --dry-run flag with --unstore. 

259 

260 Verify that with the dry-run flag the subcommand says what it would 

261 remove, but does not remove the datasets. 

262 """ 

263 self.run_test( 

264 cliArgs=["myCollection", "--dry-run", "--unstore"], 

265 exPruneDatasetsCallArgs=None, 

266 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)), 

267 exGetTablesCalled=True, 

268 exMsgs=(pruneDatasets_wouldRemoveMsg, astropyTablesToStr(getTables())), 

269 ) 

270 

271 def test_dryRun_disassociate(self): 

272 """Test the --dry-run flag with --disassociate. 

273 

274 Verify that with the dry-run flag the subcommand says what it would 

275 remove, but does not remove the datasets. 

276 """ 

277 collection = "myCollection" 

278 self.run_test( 

279 cliArgs=[collection, "--dry-run", "--disassociate", "tag1"], 

280 exPruneDatasetsCallArgs=None, 

281 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=(collection,)), 

282 exGetTablesCalled=True, 

283 exMsgs=( 

284 pruneDatasets_wouldDisassociateMsg.format(collections=(collection,)), 

285 astropyTablesToStr(getTables()), 

286 ), 

287 ) 

288 

289 def test_dryRun_unstoreAndDisassociate(self): 

290 """Test the --dry-run flag with --unstore and --disassociate. 

291 

292 Verify that with the dry-run flag the subcommand says what it would 

293 remove, but does not remove the datasets. 

294 """ 

295 collection = "myCollection" 

296 self.run_test( 

297 cliArgs=[collection, "--dry-run", "--unstore", "--disassociate", "tag1"], 

298 exPruneDatasetsCallArgs=None, 

299 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=(collection,)), 

300 exGetTablesCalled=True, 

301 exMsgs=( 

302 pruneDatasets_wouldDisassociateAndRemoveMsg.format(collections=(collection,)), 

303 astropyTablesToStr(getTables()), 

304 ), 

305 ) 

306 

307 def test_noConfirm(self): 

308 """Test the --no-confirm flag. 

309 

310 Verify that with the no-confirm flag the subcommand does not ask for 

311 a confirmation, prints the did remove message and the tables that were 

312 passed for removal. 

313 """ 

314 self.run_test( 

315 cliArgs=["myCollection", "--no-confirm", "--unstore"], 

316 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True), 

317 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)), 

318 exGetTablesCalled=True, 

319 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())), 

320 ) 

321 

322 def test_quiet(self): 

323 """Test the --quiet flag. 

324 

325 Verify that with the quiet flag and the no-confirm flags set that no 

326 output is produced by the subcommand. 

327 """ 

328 self.run_test( 

329 cliArgs=["myCollection", "--quiet", "--unstore"], 

330 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True), 

331 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)), 

332 exGetTablesCalled=True, 

333 exMsgs=None, 

334 ) 

335 

336 def test_quietWithDryRun(self): 

337 """Test for an error using the --quiet flag with --dry-run.""" 

338 self.run_test( 

339 cliArgs=["--quiet", "--dry-run", "--unstore"], 

340 exPruneDatasetsCallArgs=None, 

341 exQueryDatasetsCallArgs=None, 

342 exGetTablesCalled=False, 

343 exMsgs=(pruneDatasets_errQuietWithDryRun,), 

344 exPruneDatasetsExitCode=1, 

345 ) 

346 

347 def test_noCollections(self): 

348 """Test for an error if no collections are indicated.""" 

349 self.run_test( 

350 cliArgs=["--find-all", "--unstore"], 

351 exPruneDatasetsCallArgs=None, 

352 exQueryDatasetsCallArgs=None, 

353 exGetTablesCalled=False, 

354 exMsgs=(pruneDatasets_errNoCollectionRestriction,), 

355 exPruneDatasetsExitCode=1, 

356 ) 

357 

358 def test_noDatasets(self): 

359 """Test for expected outputs when no datasets are found.""" 

360 global doFindTables 

361 reset = doFindTables 

362 try: 

363 doFindTables = False 

364 self.run_test( 

365 cliArgs=["myCollection", "--unstore"], 

366 exPruneDatasetsCallArgs=None, 

367 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

368 repo=self.repo, collections=("myCollection",) 

369 ), 

370 exGetTablesCalled=True, 

371 exMsgs=(pruneDatasets_noDatasetsFound,), 

372 ) 

373 finally: 

374 doFindTables = reset 

375 

376 def test_purgeWithDisassociate(self): 

377 """Verify there is an error when --purge and --disassociate are both 

378 passed in. 

379 """ 

380 self.run_test( 

381 cliArgs=["--purge", "run", "--disassociate", "tag1", "tag2"], 

382 exPruneDatasetsCallArgs=None, 

383 exQueryDatasetsCallArgs=None, # should not make it far enough to call this. 

384 exGetTablesCalled=False, # ...or this. 

385 exMsgs=(pruneDatasets_errPurgeAndDisassociate,), 

386 exPruneDatasetsExitCode=1, 

387 ) 

388 

389 def test_purgeNoOp(self): 

390 """Verify there is an error when none of --purge, --unstore, or 

391 --disassociate are passed. 

392 """ 

393 self.run_test( 

394 cliArgs=[], 

395 exPruneDatasetsCallArgs=None, 

396 exQueryDatasetsCallArgs=None, # should not make it far enough to call this. 

397 exGetTablesCalled=False, # ...or this. 

398 exMsgs=(pruneDatasets_errNoOp,), 

399 exPruneDatasetsExitCode=1, 

400 ) 

401 

402 @patch.object( 402 ↛ exitline 402 didn't jump to the function exit

403 lsst.daf.butler.registry.sql_registry.SqlRegistry, 

404 "getCollectionType", 

405 side_effect=lambda x: CollectionType.RUN, 

406 ) 

407 def test_purgeImpliedArgs(self, mockGetCollectionType): 

408 """Verify the arguments implied by --purge. 

409 

410 --purge <run> implies the following arguments to butler.pruneDatasets: 

411 purge=True, disassociate=True, unstore=True 

412 And for QueryDatasets, if COLLECTIONS is not passed then <run> gets 

413 used as the value of COLLECTIONS (and when there is a COLLECTIONS 

414 value then find_first gets set to True) 

415 """ 

416 self.run_test( 

417 cliArgs=["--purge", "run"], 

418 invokeInput="yes", 

419 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

420 purge=True, refs=getDatasets(), disassociate=True, unstore=True 

421 ), 

422 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

423 repo=self.repo, collections=("run",), find_first=True 

424 ), 

425 exGetTablesCalled=True, 

426 exMsgs=( 

427 pruneDatasets_willRemoveMsg, 

428 pruneDatasets_askContinueMsg, 

429 astropyTablesToStr(getTables()), 

430 pruneDatasets_didRemoveAforementioned, 

431 ), 

432 ) 

433 

434 @patch.object( 434 ↛ exitline 434 didn't jump to the function exit

435 lsst.daf.butler.registry.sql_registry.SqlRegistry, 

436 "getCollectionType", 

437 side_effect=lambda x: CollectionType.RUN, 

438 ) 

439 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType): 

440 """Verify the arguments implied by --purge, with a COLLECTIONS.""" 

441 self.run_test( 

442 cliArgs=["myCollection", "--purge", "run"], 

443 invokeInput="yes", 

444 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

445 purge=True, disassociate=True, unstore=True, refs=getDatasets() 

446 ), 

447 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

448 repo=self.repo, collections=("myCollection",), find_first=True 

449 ), 

450 exGetTablesCalled=True, 

451 exMsgs=( 

452 pruneDatasets_willRemoveMsg, 

453 pruneDatasets_askContinueMsg, 

454 astropyTablesToStr(getTables()), 

455 pruneDatasets_didRemoveAforementioned, 

456 ), 

457 ) 

458 

459 @patch.object( 459 ↛ exitline 459 didn't jump to the function exit

460 lsst.daf.butler.registry.sql_registry.SqlRegistry, 

461 "getCollectionType", 

462 side_effect=lambda x: CollectionType.TAGGED, 

463 ) 

464 def test_purgeOnNonRunCollection(self, mockGetCollectionType): 

465 """Verify calling run on a non-run collection fails with expected 

466 error message. 

467 """ 

468 collectionName = "myTaggedCollection" 

469 self.run_test( 

470 cliArgs=["--purge", collectionName], 

471 invokeInput="yes", 

472 exPruneDatasetsCallArgs=None, 

473 exQueryDatasetsCallArgs=None, 

474 exGetTablesCalled=False, 

475 exMsgs=(pruneDatasets_errPruneOnNotRun.format(collection=collectionName)), 

476 exPruneDatasetsExitCode=1, 

477 ) 

478 

479 def test_disassociateImpliedArgs(self): 

480 """Verify the arguments implied by --disassociate. 

481 

482 --disassociate <tags> implies the following arguments to 

483 butler.pruneDatasets: 

484 disassociate=True, tags=<tags> 

485 and if COLLECTIONS is not passed then <tags> gets used as the value 

486 of COLLECTIONS. 

487 

488 Use the --no-confirm flag instead of invokeInput="yes", and check for 

489 the associated output. 

490 """ 

491 self.run_test( 

492 cliArgs=["--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"], 

493 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

494 tags=("tag1", "tag2"), disassociate=True, refs=getDatasets() 

495 ), 

496 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

497 repo=self.repo, collections=("tag1", "tag2"), find_first=True 

498 ), 

499 exGetTablesCalled=True, 

500 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())), 

501 ) 

502 

503 def test_disassociateImpliedArgsWithCollections(self): 

504 """Verify the arguments implied by --disassociate, with a --collection 

505 flag. 

506 """ 

507 self.run_test( 

508 cliArgs=["myCollection", "--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"], 

509 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

510 tags=("tag1", "tag2"), disassociate=True, refs=getDatasets() 

511 ), 

512 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

513 repo=self.repo, collections=("myCollection",), find_first=True 

514 ), 

515 exGetTablesCalled=True, 

516 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())), 

517 ) 

518 

519 

520if __name__ == "__main__": 

521 unittest.main()