Coverage for tests/test_cliCmdPruneDatasets.py: 47%

106 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-08-12 09:20 +0000

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

23""" 

24 

25import unittest 

26from unittest.mock import patch 

27 

28# Tests require the SqlRegistry 

29import lsst.daf.butler.registries.sql 

30import lsst.daf.butler.script 

31from astropy.table import Table 

32from lsst.daf.butler import Butler 

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

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

35 pruneDatasets_askContinueMsg, 

36 pruneDatasets_didNotRemoveAforementioned, 

37 pruneDatasets_didRemoveAforementioned, 

38 pruneDatasets_didRemoveMsg, 

39 pruneDatasets_errNoCollectionRestriction, 

40 pruneDatasets_errNoOp, 

41 pruneDatasets_errPruneOnNotRun, 

42 pruneDatasets_errPurgeAndDisassociate, 

43 pruneDatasets_errQuietWithDryRun, 

44 pruneDatasets_noDatasetsFound, 

45 pruneDatasets_willRemoveMsg, 

46 pruneDatasets_wouldDisassociateAndRemoveMsg, 

47 pruneDatasets_wouldDisassociateMsg, 

48 pruneDatasets_wouldRemoveMsg, 

49) 

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

51from lsst.daf.butler.registry import CollectionType 

52from lsst.daf.butler.script import QueryDatasets 

53 

54doFindTables = True 

55 

56 

57def getTables(): 

58 """Return test table.""" 

59 if doFindTables: 

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

61 return () 

62 

63 

64def getDatasets(): 

65 """Return the datasets string.""" 

66 return "datasets" 

67 

68 

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

70 """Return a query datasets object.""" 

71 return QueryDatasets(*args, **kwargs) 

72 

73 

74class PruneDatasetsTestCase(unittest.TestCase): 

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

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

77 ``scripts/_pruneDatasets.py``). 

78 

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

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

81 """ 

82 

83 def setUp(self): 

84 self.repo = "here" 

85 

86 @staticmethod 

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

88 expectedArgs = dict( 

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

90 ) 

91 expectedArgs.update(kwargs) 

92 return expectedArgs 

93 

94 @staticmethod 

95 def makePruneDatasetsArgs(**kwargs): 

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

97 expectedArgs.update(kwargs) 

98 return expectedArgs 

99 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

115 @patch.object(Butler, "pruneDatasets") 

116 def run_test( 

117 self, 

118 mockPruneDatasets, 

119 mockQueryDatasets_init, 

120 mockQueryDatasets_getDatasets, 

121 mockQueryDatasets_getTables, 

122 cliArgs, 

123 exMsgs, 

124 exPruneDatasetsCallArgs, 

125 exGetTablesCalled, 

126 exQueryDatasetsCallArgs, 

127 invokeInput=None, 

128 exPruneDatasetsExitCode=0, 

129 ): 

130 """Execute the test. 

131 

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

133 output, exit codes, and mock calls. 

134 

135 Parameters 

136 ---------- 

137 mockPruneDatasets : `MagicMock` 

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

139 mockQueryDatasets_init : `MagicMock` 

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

141 mockQueryDatasets_getDatasets : `MagicMock` 

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

143 mockQueryDatasets_getTables : `MagicMock` 

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

145 cliArgs : `list` [`str`] 

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

147 subcommand name or the repo. 

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

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

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

151 produced. 

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

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

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

155 exGetTablesCalled : bool 

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

157 `False`. 

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

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

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

161 invokeInput : `str`, optional. 

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

163 default None. 

164 exPruneDatasetsExitCode : `int` 

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

166 """ 

167 runner = LogCliRunner() 

168 with runner.isolated_filesystem(): 

169 # Make a repo so a butler can be created 

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

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

172 

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

174 # mocks: 

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

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

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

178 

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

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

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

182 # ``getDatasets()``. 

183 if exPruneDatasetsCallArgs: 

184 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs) 

185 else: 

186 mockPruneDatasets.assert_not_called() 

187 

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

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

190 # time each. 

191 if exQueryDatasetsCallArgs: 

192 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs) 

193 else: 

194 mockQueryDatasets_init.assert_not_called() 

195 # If Butler.pruneDatasets was not called, then 

196 # QueryDatasets.getDatasets also does not get called. 

197 if exPruneDatasetsCallArgs: 

198 mockQueryDatasets_getDatasets.assert_called_once() 

199 else: 

200 mockQueryDatasets_getDatasets.assert_not_called() 

201 if exGetTablesCalled: 

202 mockQueryDatasets_getTables.assert_called_once() 

203 else: 

204 mockQueryDatasets_getTables.assert_not_called() 

205 

206 if exMsgs is None: 

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

208 else: 

209 for expectedMsg in exMsgs: 

210 self.assertIn(expectedMsg, result.output) 

211 

212 def test_defaults_doContinue(self): 

213 """Test running with the default values. 

214 

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

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

217 """ 

218 self.run_test( 

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

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

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

222 exGetTablesCalled=True, 

223 exMsgs=( 

224 pruneDatasets_willRemoveMsg, 

225 pruneDatasets_askContinueMsg, 

226 astropyTablesToStr(getTables()), 

227 pruneDatasets_didRemoveAforementioned, 

228 ), 

229 invokeInput="yes", 

230 ) 

231 

232 def test_defaults_doNotContinue(self): 

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

234 

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

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

237 """ 

238 self.run_test( 

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

240 exPruneDatasetsCallArgs=None, 

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

242 exGetTablesCalled=True, 

243 exMsgs=( 

244 pruneDatasets_willRemoveMsg, 

245 pruneDatasets_askContinueMsg, 

246 pruneDatasets_didNotRemoveAforementioned, 

247 ), 

248 invokeInput="no", 

249 ) 

250 

251 def test_dryRun_unstore(self): 

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

253 

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

255 remove, but does not remove the datasets. 

256 """ 

257 self.run_test( 

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

259 exPruneDatasetsCallArgs=None, 

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

261 exGetTablesCalled=True, 

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

263 ) 

264 

265 def test_dryRun_disassociate(self): 

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

267 

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

269 remove, but does not remove the datasets. 

270 """ 

271 collection = "myCollection" 

272 self.run_test( 

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

274 exPruneDatasetsCallArgs=None, 

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

276 exGetTablesCalled=True, 

277 exMsgs=( 

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

279 astropyTablesToStr(getTables()), 

280 ), 

281 ) 

282 

283 def test_dryRun_unstoreAndDisassociate(self): 

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

285 

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

287 remove, but does not remove the datasets. 

288 """ 

289 collection = "myCollection" 

290 self.run_test( 

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

292 exPruneDatasetsCallArgs=None, 

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

294 exGetTablesCalled=True, 

295 exMsgs=( 

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

297 astropyTablesToStr(getTables()), 

298 ), 

299 ) 

300 

301 def test_noConfirm(self): 

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

303 

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

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

306 passed for removal. 

307 """ 

308 self.run_test( 

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

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

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

312 exGetTablesCalled=True, 

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

314 ) 

315 

316 def test_quiet(self): 

317 """Test the --quiet flag. 

318 

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

320 output is produced by the subcommand. 

321 """ 

322 self.run_test( 

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

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

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

326 exGetTablesCalled=True, 

327 exMsgs=None, 

328 ) 

329 

330 def test_quietWithDryRun(self): 

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

332 self.run_test( 

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

334 exPruneDatasetsCallArgs=None, 

335 exQueryDatasetsCallArgs=None, 

336 exGetTablesCalled=False, 

337 exMsgs=(pruneDatasets_errQuietWithDryRun,), 

338 exPruneDatasetsExitCode=1, 

339 ) 

340 

341 def test_noCollections(self): 

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

343 self.run_test( 

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

345 exPruneDatasetsCallArgs=None, 

346 exQueryDatasetsCallArgs=None, 

347 exGetTablesCalled=False, 

348 exMsgs=(pruneDatasets_errNoCollectionRestriction,), 

349 exPruneDatasetsExitCode=1, 

350 ) 

351 

352 def test_noDatasets(self): 

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

354 global doFindTables 

355 reset = doFindTables 

356 try: 

357 doFindTables = False 

358 self.run_test( 

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

360 exPruneDatasetsCallArgs=None, 

361 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

363 ), 

364 exGetTablesCalled=True, 

365 exMsgs=(pruneDatasets_noDatasetsFound,), 

366 ) 

367 finally: 

368 doFindTables = reset 

369 

370 def test_purgeWithDisassociate(self): 

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

372 passed in. 

373 """ 

374 self.run_test( 

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

376 exPruneDatasetsCallArgs=None, 

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

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

379 exMsgs=(pruneDatasets_errPurgeAndDisassociate,), 

380 exPruneDatasetsExitCode=1, 

381 ) 

382 

383 def test_purgeNoOp(self): 

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

385 --disassociate are passed. 

386 """ 

387 self.run_test( 

388 cliArgs=[], 

389 exPruneDatasetsCallArgs=None, 

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

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

392 exMsgs=(pruneDatasets_errNoOp,), 

393 exPruneDatasetsExitCode=1, 

394 ) 

395 

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

397 lsst.daf.butler.registries.sql.SqlRegistry, 

398 "getCollectionType", 

399 side_effect=lambda x: CollectionType.RUN, 

400 ) 

401 def test_purgeImpliedArgs(self, mockGetCollectionType): 

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

403 

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

405 purge=True, disassociate=True, unstore=True 

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

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

408 value then find_first gets set to True) 

409 """ 

410 self.run_test( 

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

412 invokeInput="yes", 

413 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

415 ), 

416 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

418 ), 

419 exGetTablesCalled=True, 

420 exMsgs=( 

421 pruneDatasets_willRemoveMsg, 

422 pruneDatasets_askContinueMsg, 

423 astropyTablesToStr(getTables()), 

424 pruneDatasets_didRemoveAforementioned, 

425 ), 

426 ) 

427 

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

429 lsst.daf.butler.registries.sql.SqlRegistry, 

430 "getCollectionType", 

431 side_effect=lambda x: CollectionType.RUN, 

432 ) 

433 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType): 

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

435 self.run_test( 

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

437 invokeInput="yes", 

438 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

440 ), 

441 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

443 ), 

444 exGetTablesCalled=True, 

445 exMsgs=( 

446 pruneDatasets_willRemoveMsg, 

447 pruneDatasets_askContinueMsg, 

448 astropyTablesToStr(getTables()), 

449 pruneDatasets_didRemoveAforementioned, 

450 ), 

451 ) 

452 

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

454 lsst.daf.butler.registries.sql.SqlRegistry, 

455 "getCollectionType", 

456 side_effect=lambda x: CollectionType.TAGGED, 

457 ) 

458 def test_purgeOnNonRunCollection(self, mockGetCollectionType): 

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

460 error message. 

461 """ 

462 collectionName = "myTaggedCollection" 

463 self.run_test( 

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

465 invokeInput="yes", 

466 exPruneDatasetsCallArgs=None, 

467 exQueryDatasetsCallArgs=None, 

468 exGetTablesCalled=False, 

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

470 exPruneDatasetsExitCode=1, 

471 ) 

472 

473 def test_disassociateImpliedArgs(self): 

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

475 

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

477 butler.pruneDatasets: 

478 disassociate=True, tags=<tags> 

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

480 of COLLECTIONS. 

481 

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

483 the associated output. 

484 """ 

485 self.run_test( 

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

487 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

489 ), 

490 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

492 ), 

493 exGetTablesCalled=True, 

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

495 ) 

496 

497 def test_disassociateImpliedArgsWithCollections(self): 

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

499 flag. 

500 """ 

501 self.run_test( 

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

503 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

505 ), 

506 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

508 ), 

509 exGetTablesCalled=True, 

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

511 ) 

512 

513 

514if __name__ == "__main__": 

515 unittest.main()