Coverage for tests/test_cliCmdPruneDatasets.py: 41%

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

106 statements  

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_errPruneOnNotRun, 

41 pruneDatasets_errPurgeAndDisassociate, 

42 pruneDatasets_errQuietWithDryRun, 

43 pruneDatasets_noDatasetsFound, 

44 pruneDatasets_willRemoveMsg, 

45 pruneDatasets_wouldDisassociateAndRemoveMsg, 

46 pruneDatasets_wouldDisassociateMsg, 

47 pruneDatasets_wouldRemoveMsg, 

48) 

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

50from lsst.daf.butler.registry import CollectionType 

51from lsst.daf.butler.script import QueryDatasets 

52 

53doFindTables = True 

54 

55 

56def getTables(): 

57 if doFindTables: 

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

59 return tuple() 

60 

61 

62def getDatasets(): 

63 return "datasets" 

64 

65 

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

67 return QueryDatasets(*args, **kwargs) 

68 

69 

70class PruneDatasetsTestCase(unittest.TestCase): 

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

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

73 ``scripts/_pruneDatasets.py``). 

74 

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

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

77 """ 

78 

79 def setUp(self): 

80 self.repo = "here" 

81 

82 @staticmethod 

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

84 expectedArgs = dict( 

85 repo=repo, collections=(), where=None, find_first=True, show_uri=False, glob=tuple() 

86 ) 

87 expectedArgs.update(kwargs) 

88 return expectedArgs 

89 

90 @staticmethod 

91 def makePruneDatasetsArgs(**kwargs): 

92 expectedArgs = dict(refs=tuple(), disassociate=False, tags=(), purge=False, run=None, unstore=False) 

93 expectedArgs.update(kwargs) 

94 return expectedArgs 

95 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

112 def run_test( 

113 self, 

114 mockPruneDatasets, 

115 mockQueryDatasets_init, 

116 mockQueryDatasets_getDatasets, 

117 mockQueryDatasets_getTables, 

118 cliArgs, 

119 exMsgs, 

120 exPruneDatasetsCallArgs, 

121 exGetTablesCalled, 

122 exQueryDatasetsCallArgs, 

123 invokeInput=None, 

124 exPruneDatasetsExitCode=0, 

125 ): 

126 """Execute the test. 

127 

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

129 output, exit codes, and mock calls. 

130 

131 Parameters 

132 ---------- 

133 mockPruneDatasets : `MagicMock` 

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

135 mockQueryDatasets_init : `MagicMock` 

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

137 mockQueryDatasets_getDatasets : `MagicMock` 

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

139 mockQueryDatasets_getTables : `MagicMock` 

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

141 cliArgs : `list` [`str`] 

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

143 subcommand name or the repo. 

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

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

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

147 produced. 

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

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

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

151 exGetTablesCalled : bool 

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

153 `False`. 

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

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

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

157 invokeInput : `str`, optional. 

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

159 default None. 

160 exPruneDatasetsExitCode : `int` 

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

162 """ 

163 runner = LogCliRunner() 

164 with runner.isolated_filesystem(): 

165 # Make a repo so a butler can be created 

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

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

168 

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

170 # mocks: 

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

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

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

174 

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

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

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

178 # ``getDatasets()``. 

179 if exPruneDatasetsCallArgs: 

180 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs) 

181 else: 

182 mockPruneDatasets.assert_not_called() 

183 

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

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

186 # time each. 

187 if exQueryDatasetsCallArgs: 

188 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs) 

189 else: 

190 mockQueryDatasets_init.assert_not_called() 

191 # If Butler.pruneDatasets was not called, then 

192 # QueryDatasets.getDatasets also does not get called. 

193 if exPruneDatasetsCallArgs: 

194 mockQueryDatasets_getDatasets.assert_called_once() 

195 else: 

196 mockQueryDatasets_getDatasets.assert_not_called() 

197 if exGetTablesCalled: 

198 mockQueryDatasets_getTables.assert_called_once() 

199 else: 

200 mockQueryDatasets_getTables.assert_not_called() 

201 

202 if exMsgs is None: 

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

204 else: 

205 for expectedMsg in exMsgs: 

206 self.assertIn(expectedMsg, result.output) 

207 

208 def test_defaults_doContinue(self): 

209 """Test running with the default values. 

210 

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

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

213 self.run_test( 

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

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

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

217 exGetTablesCalled=True, 

218 exMsgs=( 

219 pruneDatasets_willRemoveMsg, 

220 pruneDatasets_askContinueMsg, 

221 astropyTablesToStr(getTables()), 

222 pruneDatasets_didRemoveAforementioned, 

223 ), 

224 invokeInput="yes", 

225 ) 

226 

227 def test_defaults_doNotContinue(self): 

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

229 

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

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

232 self.run_test( 

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

234 exPruneDatasetsCallArgs=None, 

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

236 exGetTablesCalled=True, 

237 exMsgs=( 

238 pruneDatasets_willRemoveMsg, 

239 pruneDatasets_askContinueMsg, 

240 pruneDatasets_didNotRemoveAforementioned, 

241 ), 

242 invokeInput="no", 

243 ) 

244 

245 def test_dryRun_unstore(self): 

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

247 

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

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

250 self.run_test( 

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

252 exPruneDatasetsCallArgs=None, 

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

254 exGetTablesCalled=True, 

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

256 ) 

257 

258 def test_dryRun_disassociate(self): 

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

260 

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

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

263 collection = "myCollection" 

264 self.run_test( 

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

266 exPruneDatasetsCallArgs=None, 

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

268 exGetTablesCalled=True, 

269 exMsgs=( 

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

271 astropyTablesToStr(getTables()), 

272 ), 

273 ) 

274 

275 def test_dryRun_unstoreAndDisassociate(self): 

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

277 

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

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

280 collection = "myCollection" 

281 self.run_test( 

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

283 exPruneDatasetsCallArgs=None, 

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

285 exGetTablesCalled=True, 

286 exMsgs=( 

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

288 astropyTablesToStr(getTables()), 

289 ), 

290 ) 

291 

292 def test_noConfirm(self): 

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

294 

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

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

297 passed for removal.""" 

298 self.run_test( 

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

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

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

302 exGetTablesCalled=True, 

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

304 ) 

305 

306 def test_quiet(self): 

307 """Test the --quiet flag. 

308 

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

310 output is produced by the subcommand.""" 

311 self.run_test( 

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

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

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

315 exGetTablesCalled=True, 

316 exMsgs=None, 

317 ) 

318 

319 def test_quietWithDryRun(self): 

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

321 self.run_test( 

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

323 exPruneDatasetsCallArgs=None, 

324 exQueryDatasetsCallArgs=None, 

325 exGetTablesCalled=False, 

326 exMsgs=(pruneDatasets_errQuietWithDryRun,), 

327 exPruneDatasetsExitCode=1, 

328 ) 

329 

330 def test_noCollections(self): 

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

332 self.run_test( 

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

334 exPruneDatasetsCallArgs=None, 

335 exQueryDatasetsCallArgs=None, 

336 exGetTablesCalled=False, 

337 exMsgs=(pruneDatasets_errNoCollectionRestriction,), 

338 exPruneDatasetsExitCode=1, 

339 ) 

340 

341 def test_noDatasets(self): 

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

343 global doFindTables 

344 reset = doFindTables 

345 try: 

346 doFindTables = False 

347 self.run_test( 

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

349 exPruneDatasetsCallArgs=None, 

350 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

352 ), 

353 exGetTablesCalled=True, 

354 exMsgs=(pruneDatasets_noDatasetsFound,), 

355 ) 

356 finally: 

357 doFindTables = reset 

358 

359 def test_purgeWithDisassociate(self): 

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

361 passed in.""" 

362 self.run_test( 

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

364 exPruneDatasetsCallArgs=None, 

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

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

367 exMsgs=(pruneDatasets_errPurgeAndDisassociate,), 

368 exPruneDatasetsExitCode=1, 

369 ) 

370 

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

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

373 "getCollectionType", 

374 side_effect=lambda x: CollectionType.RUN, 

375 ) 

376 def test_purgeImpliedArgs(self, mockGetCollectionType): 

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

378 

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

380 purge=True, run=<run>, disassociate=True, unstore=True 

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

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

383 value then find_first gets set to True) 

384 """ 

385 self.run_test( 

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

387 invokeInput="yes", 

388 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

389 purge=True, run="run", refs=getDatasets(), disassociate=True, unstore=True 

390 ), 

391 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

393 ), 

394 exGetTablesCalled=True, 

395 exMsgs=( 

396 pruneDatasets_willRemoveMsg, 

397 pruneDatasets_askContinueMsg, 

398 astropyTablesToStr(getTables()), 

399 pruneDatasets_didRemoveAforementioned, 

400 ), 

401 ) 

402 

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

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

405 "getCollectionType", 

406 side_effect=lambda x: CollectionType.RUN, 

407 ) 

408 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType): 

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

410 self.run_test( 

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

412 invokeInput="yes", 

413 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

415 ), 

416 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

417 repo=self.repo, collections=("myCollection",), 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.TAGGED, 

432 ) 

433 def test_purgeOnNonRunCollection(self, mockGetCollectionType): 

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

435 error message.""" 

436 collectionName = "myTaggedCollection" 

437 self.run_test( 

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

439 invokeInput="yes", 

440 exPruneDatasetsCallArgs=None, 

441 exQueryDatasetsCallArgs=None, 

442 exGetTablesCalled=False, 

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

444 exPruneDatasetsExitCode=1, 

445 ) 

446 

447 def test_disassociateImpliedArgs(self): 

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

449 

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

451 butler.pruneDatasets: 

452 disassociate=True, tags=<tags> 

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

454 of COLLECTIONS. 

455 

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

457 the associated output. 

458 """ 

459 self.run_test( 

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

461 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

463 ), 

464 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

466 ), 

467 exGetTablesCalled=True, 

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

469 ) 

470 

471 def test_disassociateImpliedArgsWithCollections(self): 

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

473 flag.""" 

474 self.run_test( 

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

476 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

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

478 ), 

479 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs( 

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

481 ), 

482 exGetTablesCalled=True, 

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

484 ) 

485 

486 

487if __name__ == "__main__": 487 ↛ 488line 487 didn't jump to line 488, because the condition on line 487 was never true

488 unittest.main()