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
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
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/>.
22"""Unit tests for daf_butler CLI prune-datasets subcommand.
23"""
25import unittest
26from unittest.mock import patch
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
53doFindTables = True
56def getTables():
57 if doFindTables:
58 return (Table(((1, 2, 3),), names=("foo",)),)
59 return tuple()
62def getDatasets():
63 return "datasets"
66def makeQueryDatasets(*args, **kwargs):
67 return QueryDatasets(*args, **kwargs)
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``).
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 """
79 def setUp(self):
80 self.repo = "here"
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
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
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.
128 Makes a temporary repo, invokes ``prune-datasets``. Verifies expected
129 output, exit codes, and mock calls.
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))
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))
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()
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()
202 if exMsgs is None:
203 self.assertEqual("", result.output)
204 else:
205 for expectedMsg in exMsgs:
206 self.assertIn(expectedMsg, result.output)
208 def test_defaults_doContinue(self):
209 """Test running with the default values.
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 )
227 def test_defaults_doNotContinue(self):
228 """Test running with the default values but not continuing.
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 )
245 def test_dryRun_unstore(self):
246 """Test the --dry-run flag with --unstore.
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 )
258 def test_dryRun_disassociate(self):
259 """Test the --dry-run flag with --disassociate.
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 )
275 def test_dryRun_unstoreAndDisassociate(self):
276 """Test the --dry-run flag with --unstore and --disassociate.
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 )
292 def test_noConfirm(self):
293 """Test the --no-confirm flag.
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 )
306 def test_quiet(self):
307 """Test the --quiet flag.
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 )
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 )
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 )
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
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 )
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.
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 )
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 )
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 )
447 def test_disassociateImpliedArgs(self):
448 """Verify the arguments implied by --disassociate.
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.
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 )
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 )
487if __name__ == "__main__": 487 ↛ 488line 487 didn't jump to line 488, because the condition on line 487 was never true
488 unittest.main()