Hide keyboard shortcuts

Hot-keys 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

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 

25from astropy.table import Table 

26import unittest 

27from unittest.mock import patch 

28 

29from lsst.daf.butler import Butler 

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

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

32 pruneDatasets_wouldRemoveMsg, 

33 pruneDatasets_wouldDisassociateMsg, 

34 pruneDatasets_wouldDisassociateAndRemoveMsg, 

35 pruneDatasets_willRemoveMsg, 

36 pruneDatasets_askContinueMsg, 

37 pruneDatasets_didRemoveAforementioned, 

38 pruneDatasets_didNotRemoveAforementioned, 

39 pruneDatasets_didRemoveMsg, 

40 pruneDatasets_noDatasetsFound, 

41 pruneDatasets_errPurgeAndDisassociate, 

42 pruneDatasets_errQuietWithDryRun, 

43 pruneDatasets_errNoCollectionRestriction, 

44 pruneDatasets_errPruneOnNotRun, 

45) 

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

47from lsst.daf.butler.registry import CollectionType 

48import lsst.daf.butler.script 

49from lsst.daf.butler.script import QueryDatasets 

50 

51 

52doFindTables = True 

53 

54 

55def getTables(): 

56 if doFindTables: 

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

58 return tuple() 

59 

60 

61def getDatasets(): 

62 return "datasets" 

63 

64 

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

66 return QueryDatasets(*args, **kwargs) 

67 

68 

69class PruneDatasetsTestCase(unittest.TestCase): 

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

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

72 ``scripts/_pruneDatasets.py``). 

73 

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

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

76 """ 

77 

78 def setUp(self): 

79 self.repo = "here" 

80 

81 @staticmethod 

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

83 expectedArgs = dict(repo=repo, collections=(), where=None, find_first=True, show_uri=False, 

84 glob=tuple()) 

85 expectedArgs.update(kwargs) 

86 return expectedArgs 

87 

88 @staticmethod 

89 def makePruneDatasetsArgs(**kwargs): 

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

91 unstore=False) 

92 expectedArgs.update(kwargs) 

93 return expectedArgs 

94 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

111 def run_test(self, 

112 mockPruneDatasets, 

113 mockQueryDatasets_init, 

114 mockQueryDatasets_getDatasets, 

115 mockQueryDatasets_getTables, 

116 cliArgs, 

117 exMsgs, 

118 exPruneDatasetsCallArgs, 

119 exGetTablesCalled, 

120 exQueryDatasetsCallArgs, 

121 invokeInput=None, 

122 exPruneDatasetsExitCode=0): 

123 """Execute the test. 

124 

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

126 output, exit codes, and mock calls. 

127 

128 Parameters 

129 ---------- 

130 mockPruneDatasets : `MagicMock` 

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

132 mockQueryDatasets_init : `MagicMock` 

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

134 mockQueryDatasets_getDatasets : `MagicMock` 

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

136 mockQueryDatasets_getTables : `MagicMock` 

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

138 cliArgs : `list` [`str`] 

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

140 subcommand name or the repo. 

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

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

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

144 produced. 

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

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

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

148 exGetTablesCalled : bool 

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

150 `False`. 

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

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

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

154 invokeInput : `str`, optional. 

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

156 default None. 

157 exPruneDatasetsExitCode : `int` 

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

159 """ 

160 runner = LogCliRunner() 

161 with runner.isolated_filesystem(): 

162 # Make a repo so a butler can be created 

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

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

165 

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

167 # mocks: 

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

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

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

171 

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

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

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

175 # ``getDatasets()``. 

176 if exPruneDatasetsCallArgs: 

177 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs) 

178 else: 

179 mockPruneDatasets.assert_not_called() 

180 

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

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

183 # time each. 

184 if exQueryDatasetsCallArgs: 

185 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs) 

186 else: 

187 mockQueryDatasets_init.assert_not_called() 

188 # If Butler.pruneDatasets was not called, then 

189 # QueryDatasets.getDatasets also does not get called. 

190 if exPruneDatasetsCallArgs: 

191 mockQueryDatasets_getDatasets.assert_called_once() 

192 else: 

193 mockQueryDatasets_getDatasets.assert_not_called() 

194 if exGetTablesCalled: 

195 mockQueryDatasets_getTables.assert_called_once() 

196 else: 

197 mockQueryDatasets_getTables.assert_not_called() 

198 

199 if exMsgs is None: 

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

201 else: 

202 for expectedMsg in exMsgs: 

203 self.assertIn(expectedMsg, result.output) 

204 

205 def test_defaults_doContinue(self): 

206 """Test running with the default values. 

207 

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

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

210 self.run_test(cliArgs=["myCollection", "--unstore"], 

211 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), 

212 unstore=True), 

213 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

214 collections=("myCollection",)), 

215 exGetTablesCalled=True, 

216 exMsgs=(pruneDatasets_willRemoveMsg, 

217 pruneDatasets_askContinueMsg, 

218 astropyTablesToStr(getTables()), 

219 pruneDatasets_didRemoveAforementioned), 

220 invokeInput="yes") 

221 

222 def test_defaults_doNotContinue(self): 

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

224 

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

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

227 self.run_test(cliArgs=["myCollection", "--unstore"], 

228 exPruneDatasetsCallArgs=None, 

229 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

230 collections=("myCollection",)), 

231 exGetTablesCalled=True, 

232 exMsgs=(pruneDatasets_willRemoveMsg, 

233 pruneDatasets_askContinueMsg, 

234 pruneDatasets_didNotRemoveAforementioned), 

235 invokeInput="no") 

236 

237 def test_dryRun_unstore(self): 

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

239 

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

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

242 self.run_test(cliArgs=["myCollection", "--dry-run", "--unstore"], 

243 exPruneDatasetsCallArgs=None, 

244 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

245 collections=("myCollection",)), 

246 exGetTablesCalled=True, 

247 exMsgs=(pruneDatasets_wouldRemoveMsg, 

248 astropyTablesToStr(getTables()))) 

249 

250 def test_dryRun_disassociate(self): 

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

252 

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

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

255 collection = "myCollection" 

256 self.run_test(cliArgs=[collection, "--dry-run", "--disassociate", "tag1"], 

257 exPruneDatasetsCallArgs=None, 

258 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

259 collections=(collection,)), 

260 exGetTablesCalled=True, 

261 exMsgs=(pruneDatasets_wouldDisassociateMsg.format(collections=(collection,)), 

262 astropyTablesToStr(getTables()))) 

263 

264 def test_dryRun_unstoreAndDisassociate(self): 

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

266 

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

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

269 collection = "myCollection" 

270 self.run_test(cliArgs=[collection, "--dry-run", "--unstore", "--disassociate", "tag1"], 

271 exPruneDatasetsCallArgs=None, 

272 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

273 collections=(collection,)), 

274 exGetTablesCalled=True, 

275 exMsgs=(pruneDatasets_wouldDisassociateAndRemoveMsg.format(collections=(collection,)), 

276 astropyTablesToStr(getTables()))) 

277 

278 def test_noConfirm(self): 

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

280 

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

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

283 passed for removal.""" 

284 self.run_test(cliArgs=["myCollection", "--no-confirm", "--unstore"], 

285 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), 

286 unstore=True), 

287 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

288 collections=("myCollection",)), 

289 exGetTablesCalled=True, 

290 exMsgs=(pruneDatasets_didRemoveMsg, 

291 astropyTablesToStr(getTables()))) 

292 

293 def test_quiet(self): 

294 """Test the --quiet flag. 

295 

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

297 output is produced by the subcommand.""" 

298 self.run_test(cliArgs=["myCollection", "--quiet", "--unstore"], 

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

300 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

301 collections=("myCollection",)), 

302 exGetTablesCalled=True, 

303 exMsgs=None) 

304 

305 def test_quietWithDryRun(self): 

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

307 """ 

308 self.run_test(cliArgs=["--quiet", "--dry-run", "--unstore"], 

309 exPruneDatasetsCallArgs=None, 

310 exQueryDatasetsCallArgs=None, 

311 exGetTablesCalled=False, 

312 exMsgs=(pruneDatasets_errQuietWithDryRun,), 

313 exPruneDatasetsExitCode=1) 

314 

315 def test_noCollections(self): 

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

317 """ 

318 self.run_test(cliArgs=["--find-all", "--unstore"], 

319 exPruneDatasetsCallArgs=None, 

320 exQueryDatasetsCallArgs=None, 

321 exGetTablesCalled=False, 

322 exMsgs=(pruneDatasets_errNoCollectionRestriction,), 

323 exPruneDatasetsExitCode=1) 

324 

325 def test_noDatasets(self): 

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

327 global doFindTables 

328 reset = doFindTables 

329 try: 

330 doFindTables = False 

331 self.run_test(cliArgs=["myCollection", "--unstore"], 

332 exPruneDatasetsCallArgs=None, 

333 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

334 collections=("myCollection",)), 

335 exGetTablesCalled=True, 

336 exMsgs=(pruneDatasets_noDatasetsFound,)) 

337 finally: 

338 doFindTables = reset 

339 

340 def test_purgeWithDisassociate(self): 

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

342 passed in. """ 

343 self.run_test( 

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

345 exPruneDatasetsCallArgs=None, 

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

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

348 exMsgs=(pruneDatasets_errPurgeAndDisassociate,), 

349 exPruneDatasetsExitCode=1 

350 ) 

351 

352 @patch.object(lsst.daf.butler.Registry, "getCollectionType", 352 ↛ exitline 352 didn't jump to the function exit

353 side_effect=lambda x: CollectionType.RUN) 

354 def test_purgeImpliedArgs(self, mockGetCollectionType): 

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

356 

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

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

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

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

361 value then find_first gets set to True) 

362 """ 

363 self.run_test( 

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

365 invokeInput="yes", 

366 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

367 purge=True, 

368 run="run", 

369 refs=getDatasets(), 

370 disassociate=True, 

371 unstore=True 

372 ), 

373 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

374 collections=("run",), 

375 find_first=True), 

376 exGetTablesCalled=True, 

377 exMsgs=(pruneDatasets_willRemoveMsg, 

378 pruneDatasets_askContinueMsg, 

379 astropyTablesToStr(getTables()), 

380 pruneDatasets_didRemoveAforementioned) 

381 ) 

382 

383 @patch.object(lsst.daf.butler.Registry, "getCollectionType", 383 ↛ exitline 383 didn't jump to the function exit

384 side_effect=lambda x: CollectionType.RUN) 

385 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType): 

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

387 self.run_test( 

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

389 invokeInput="yes", 

390 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

391 purge=True, 

392 run="run", 

393 disassociate=True, 

394 unstore=True, 

395 refs=getDatasets() 

396 ), 

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

398 find_first=True), 

399 exGetTablesCalled=True, 

400 exMsgs=(pruneDatasets_willRemoveMsg, 

401 pruneDatasets_askContinueMsg, 

402 astropyTablesToStr(getTables()), 

403 pruneDatasets_didRemoveAforementioned) 

404 ) 

405 

406 @patch.object(lsst.daf.butler.Registry, "getCollectionType", 406 ↛ exitline 406 didn't jump to the function exit

407 side_effect=lambda x: CollectionType.TAGGED) 

408 def test_purgeOnNonRunCollection(self, mockGetCollectionType): 

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

410 error message. """ 

411 collectionName = "myTaggedCollection" 

412 self.run_test( 

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

414 invokeInput="yes", 

415 exPruneDatasetsCallArgs=None, 

416 exQueryDatasetsCallArgs=None, 

417 exGetTablesCalled=False, 

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

419 exPruneDatasetsExitCode=1, 

420 ) 

421 

422 def test_disassociateImpliedArgs(self): 

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

424 

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

426 butler.pruneDatasets: 

427 disassociate=True, tags=<tags> 

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

429 of COLLETIONS. 

430 

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

432 the associated output. 

433 """ 

434 self.run_test( 

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

436 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

437 tags=("tag1", "tag2"), 

438 disassociate=True, 

439 refs=getDatasets() 

440 ), 

441 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

442 collections=("tag1", "tag2"), 

443 find_first=True), 

444 exGetTablesCalled=True, 

445 exMsgs=(pruneDatasets_didRemoveMsg, 

446 astropyTablesToStr(getTables())) 

447 ) 

448 

449 def test_disassociateImpliedArgsWithCollections(self): 

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

451 flag.""" 

452 self.run_test( 

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

454 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs( 

455 tags=("tag1", "tag2"), 

456 disassociate=True, 

457 refs=getDatasets() 

458 ), 

459 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, 

460 collections=("myCollection",), 

461 find_first=True), 

462 exGetTablesCalled=True, 

463 exMsgs=(pruneDatasets_didRemoveMsg, 

464 astropyTablesToStr(getTables())) 

465 ) 

466 

467 

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

469 unittest.main()