Coverage for tests/test_cliCmdPruneDatasets.py: 39%

106 statements  

« prev     ^ index     » next       coverage.py v7.2.5, created at 2023-05-03 09:15 +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 if doFindTables: 

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

60 return tuple() 

61 

62 

63def getDatasets(): 

64 return "datasets" 

65 

66 

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

68 return QueryDatasets(*args, **kwargs) 

69 

70 

71class PruneDatasetsTestCase(unittest.TestCase): 

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

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

74 ``scripts/_pruneDatasets.py``). 

75 

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

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

78 """ 

79 

80 def setUp(self): 

81 self.repo = "here" 

82 

83 @staticmethod 

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

85 expectedArgs = dict( 

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

87 ) 

88 expectedArgs.update(kwargs) 

89 return expectedArgs 

90 

91 @staticmethod 

92 def makePruneDatasetsArgs(**kwargs): 

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

94 expectedArgs.update(kwargs) 

95 return expectedArgs 

96 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

113 def run_test( 

114 self, 

115 mockPruneDatasets, 

116 mockQueryDatasets_init, 

117 mockQueryDatasets_getDatasets, 

118 mockQueryDatasets_getTables, 

119 cliArgs, 

120 exMsgs, 

121 exPruneDatasetsCallArgs, 

122 exGetTablesCalled, 

123 exQueryDatasetsCallArgs, 

124 invokeInput=None, 

125 exPruneDatasetsExitCode=0, 

126 ): 

127 """Execute the test. 

128 

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

130 output, exit codes, and mock calls. 

131 

132 Parameters 

133 ---------- 

134 mockPruneDatasets : `MagicMock` 

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

136 mockQueryDatasets_init : `MagicMock` 

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

138 mockQueryDatasets_getDatasets : `MagicMock` 

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

140 mockQueryDatasets_getTables : `MagicMock` 

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

142 cliArgs : `list` [`str`] 

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

144 subcommand name or the repo. 

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

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

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

148 produced. 

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

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

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

152 exGetTablesCalled : bool 

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

154 `False`. 

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

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

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

158 invokeInput : `str`, optional. 

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

160 default None. 

161 exPruneDatasetsExitCode : `int` 

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

163 """ 

164 runner = LogCliRunner() 

165 with runner.isolated_filesystem(): 

166 # Make a repo so a butler can be created 

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

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

169 

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

171 # mocks: 

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

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

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

175 

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

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

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

179 # ``getDatasets()``. 

180 if exPruneDatasetsCallArgs: 

181 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs) 

182 else: 

183 mockPruneDatasets.assert_not_called() 

184 

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

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

187 # time each. 

188 if exQueryDatasetsCallArgs: 

189 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs) 

190 else: 

191 mockQueryDatasets_init.assert_not_called() 

192 # If Butler.pruneDatasets was not called, then 

193 # QueryDatasets.getDatasets also does not get called. 

194 if exPruneDatasetsCallArgs: 

195 mockQueryDatasets_getDatasets.assert_called_once() 

196 else: 

197 mockQueryDatasets_getDatasets.assert_not_called() 

198 if exGetTablesCalled: 

199 mockQueryDatasets_getTables.assert_called_once() 

200 else: 

201 mockQueryDatasets_getTables.assert_not_called() 

202 

203 if exMsgs is None: 

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

205 else: 

206 for expectedMsg in exMsgs: 

207 self.assertIn(expectedMsg, result.output) 

208 

209 def test_defaults_doContinue(self): 

210 """Test running with the default values. 

211 

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

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

214 self.run_test( 

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

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

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

218 exGetTablesCalled=True, 

219 exMsgs=( 

220 pruneDatasets_willRemoveMsg, 

221 pruneDatasets_askContinueMsg, 

222 astropyTablesToStr(getTables()), 

223 pruneDatasets_didRemoveAforementioned, 

224 ), 

225 invokeInput="yes", 

226 ) 

227 

228 def test_defaults_doNotContinue(self): 

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

230 

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

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

233 self.run_test( 

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

235 exPruneDatasetsCallArgs=None, 

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

237 exGetTablesCalled=True, 

238 exMsgs=( 

239 pruneDatasets_willRemoveMsg, 

240 pruneDatasets_askContinueMsg, 

241 pruneDatasets_didNotRemoveAforementioned, 

242 ), 

243 invokeInput="no", 

244 ) 

245 

246 def test_dryRun_unstore(self): 

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

248 

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

250 remove, but does not remove the datasets.""" 

251 self.run_test( 

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

253 exPruneDatasetsCallArgs=None, 

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

255 exGetTablesCalled=True, 

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

257 ) 

258 

259 def test_dryRun_disassociate(self): 

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

261 

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

263 remove, but does not remove the datasets.""" 

264 collection = "myCollection" 

265 self.run_test( 

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

267 exPruneDatasetsCallArgs=None, 

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

269 exGetTablesCalled=True, 

270 exMsgs=( 

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

272 astropyTablesToStr(getTables()), 

273 ), 

274 ) 

275 

276 def test_dryRun_unstoreAndDisassociate(self): 

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

278 

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

280 remove, but does not remove the datasets.""" 

281 collection = "myCollection" 

282 self.run_test( 

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

284 exPruneDatasetsCallArgs=None, 

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

286 exGetTablesCalled=True, 

287 exMsgs=( 

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

289 astropyTablesToStr(getTables()), 

290 ), 

291 ) 

292 

293 def test_noConfirm(self): 

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

295 

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

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

298 passed for removal.""" 

299 self.run_test( 

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

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

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

303 exGetTablesCalled=True, 

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

305 ) 

306 

307 def test_quiet(self): 

308 """Test the --quiet flag. 

309 

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

311 output is produced by the subcommand.""" 

312 self.run_test( 

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

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

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

316 exGetTablesCalled=True, 

317 exMsgs=None, 

318 ) 

319 

320 def test_quietWithDryRun(self): 

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

322 self.run_test( 

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

324 exPruneDatasetsCallArgs=None, 

325 exQueryDatasetsCallArgs=None, 

326 exGetTablesCalled=False, 

327 exMsgs=(pruneDatasets_errQuietWithDryRun,), 

328 exPruneDatasetsExitCode=1, 

329 ) 

330 

331 def test_noCollections(self): 

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

333 self.run_test( 

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

335 exPruneDatasetsCallArgs=None, 

336 exQueryDatasetsCallArgs=None, 

337 exGetTablesCalled=False, 

338 exMsgs=(pruneDatasets_errNoCollectionRestriction,), 

339 exPruneDatasetsExitCode=1, 

340 ) 

341 

342 def test_noDatasets(self): 

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

344 global doFindTables 

345 reset = doFindTables 

346 try: 

347 doFindTables = False 

348 self.run_test( 

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

350 exPruneDatasetsCallArgs=None, 

351 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

353 ), 

354 exGetTablesCalled=True, 

355 exMsgs=(pruneDatasets_noDatasetsFound,), 

356 ) 

357 finally: 

358 doFindTables = reset 

359 

360 def test_purgeWithDisassociate(self): 

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

362 passed in.""" 

363 self.run_test( 

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

365 exPruneDatasetsCallArgs=None, 

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

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

368 exMsgs=(pruneDatasets_errPurgeAndDisassociate,), 

369 exPruneDatasetsExitCode=1, 

370 ) 

371 

372 def test_purgeNoOp(self): 

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

374 --disassociate are passed.""" 

375 self.run_test( 

376 cliArgs=[], 

377 exPruneDatasetsCallArgs=None, 

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

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

380 exMsgs=(pruneDatasets_errNoOp,), 

381 exPruneDatasetsExitCode=1, 

382 ) 

383 

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

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

386 "getCollectionType", 

387 side_effect=lambda x: CollectionType.RUN, 

388 ) 

389 def test_purgeImpliedArgs(self, mockGetCollectionType): 

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

391 

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

393 purge=True, disassociate=True, unstore=True 

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

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

396 value then find_first gets set to True) 

397 """ 

398 self.run_test( 

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

400 invokeInput="yes", 

401 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

403 ), 

404 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

406 ), 

407 exGetTablesCalled=True, 

408 exMsgs=( 

409 pruneDatasets_willRemoveMsg, 

410 pruneDatasets_askContinueMsg, 

411 astropyTablesToStr(getTables()), 

412 pruneDatasets_didRemoveAforementioned, 

413 ), 

414 ) 

415 

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

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

418 "getCollectionType", 

419 side_effect=lambda x: CollectionType.RUN, 

420 ) 

421 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType): 

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

423 self.run_test( 

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

425 invokeInput="yes", 

426 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

428 ), 

429 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

431 ), 

432 exGetTablesCalled=True, 

433 exMsgs=( 

434 pruneDatasets_willRemoveMsg, 

435 pruneDatasets_askContinueMsg, 

436 astropyTablesToStr(getTables()), 

437 pruneDatasets_didRemoveAforementioned, 

438 ), 

439 ) 

440 

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

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

443 "getCollectionType", 

444 side_effect=lambda x: CollectionType.TAGGED, 

445 ) 

446 def test_purgeOnNonRunCollection(self, mockGetCollectionType): 

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

448 error message.""" 

449 collectionName = "myTaggedCollection" 

450 self.run_test( 

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

452 invokeInput="yes", 

453 exPruneDatasetsCallArgs=None, 

454 exQueryDatasetsCallArgs=None, 

455 exGetTablesCalled=False, 

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

457 exPruneDatasetsExitCode=1, 

458 ) 

459 

460 def test_disassociateImpliedArgs(self): 

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

462 

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

464 butler.pruneDatasets: 

465 disassociate=True, tags=<tags> 

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

467 of COLLECTIONS. 

468 

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

470 the associated output. 

471 """ 

472 self.run_test( 

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

474 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

476 ), 

477 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

479 ), 

480 exGetTablesCalled=True, 

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

482 ) 

483 

484 def test_disassociateImpliedArgsWithCollections(self): 

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

486 flag.""" 

487 self.run_test( 

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

489 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

491 ), 

492 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

494 ), 

495 exGetTablesCalled=True, 

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

497 ) 

498 

499 

500if __name__ == "__main__": 

501 unittest.main()