Coverage for tests/test_cliCmdPruneDatasets.py: 38%
110 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-25 00:34 -0800
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28"""Unit tests for daf_butler CLI prune-datasets subcommand."""
30import unittest
31from itertools import chain
32from unittest.mock import ANY, patch
34from astropy.table import Table
36# Tests require the SqlRegistry
37import lsst.daf.butler.registry.sql_registry
38import lsst.daf.butler.script
39from lsst.daf.butler import CollectionInfo, CollectionType
40from lsst.daf.butler.cli.butler import cli as butlerCli
41from lsst.daf.butler.cli.cmd.commands import (
42 pruneDatasets_askContinueMsg,
43 pruneDatasets_didNotRemoveAforementioned,
44 pruneDatasets_didRemoveAforementioned,
45 pruneDatasets_didRemoveMsg,
46 pruneDatasets_errNoCollectionRestriction,
47 pruneDatasets_errNoOp,
48 pruneDatasets_errPruneOnNotRun,
49 pruneDatasets_errPurgeAndDisassociate,
50 pruneDatasets_errQuietWithDryRun,
51 pruneDatasets_noDatasetsFound,
52 pruneDatasets_willRemoveMsg,
53 pruneDatasets_wouldDisassociateAndRemoveMsg,
54 pruneDatasets_wouldDisassociateMsg,
55 pruneDatasets_wouldRemoveMsg,
56)
57from lsst.daf.butler.cli.utils import LogCliRunner, astropyTablesToStr, clickResultMsg
58from lsst.daf.butler.direct_butler import DirectButler
59from lsst.daf.butler.script import QueryDatasets
61doFindTables = True
64def getTables():
65 """Return test table."""
66 if doFindTables:
67 yield from (Table(((1, 2, 3),), names=("foo",)),)
68 return ()
71def getDatasets():
72 """Return the datasets string."""
73 yield ["datasets"]
76def getRefs():
77 """Return all the datasets as a single list."""
78 return list(chain(*getDatasets()))
81def makeQueryDatasets(*args, **kwargs):
82 """Return a query datasets object."""
83 if not kwargs.get("glob"):
84 # Use all dataset types if not specified.
85 kwargs["glob"] = ("*",)
86 return QueryDatasets(*args, **kwargs)
89class PruneDatasetsTestCase(unittest.TestCase):
90 """Tests the ``prune_datasets`` "command" function (in
91 ``cli/cmd/commands.py``) and the ``pruneDatasets`` "script" function (in
92 ``scripts/_pruneDatasets.py``).
94 ``Butler.pruneDatasets`` and a few other functions that get called before
95 it are mocked, and tests check for expected arguments to those mocks.
96 """
98 def setUp(self):
99 self.repo = "here"
101 @staticmethod
102 def makeQueryDatasetsArgs(**kwargs):
103 expectedArgs = dict(
104 butler=ANY, collections=("*",), where="", find_first=True, show_uri=False, glob=tuple()
105 )
106 expectedArgs.update(kwargs)
107 return expectedArgs
109 @staticmethod
110 def makePruneDatasetsArgs(**kwargs):
111 expectedArgs = dict(refs=tuple(), disassociate=False, tags=(), purge=False, unstore=False)
112 expectedArgs.update(kwargs)
113 return expectedArgs
115 # Mock the QueryDatasets.getTables function to return a set of Astropy
116 # tables, similar to what would be returned by a call to
117 # QueryDatasets.getTables on a repo with real data.
118 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getTables", side_effect=getTables)
119 # Mock the QueryDatasets.getDatasets function. Normally it would return a
120 # list of queries.DatasetQueryResults, but all we need to do is verify that
121 # the output of this function is passed into our pruneDatasets magicMock,
122 # so we can return something arbitrary that we can test is equal."""
123 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getDatasets", side_effect=getDatasets)
124 # Mock the actual QueryDatasets class, so we can inspect calls to its init
125 # function. Note that the side_effect returns an instance of QueryDatasets,
126 # so this mock records and then is a pass-through.
127 @patch.object(lsst.daf.butler.script._pruneDatasets, "QueryDatasets", side_effect=makeQueryDatasets)
128 # Mock the pruneDatasets butler command so we can test for expected calls
129 # to it, without dealing with setting up a full repo with data for it.
130 @patch.object(DirectButler, "pruneDatasets")
131 def run_test(
132 self,
133 mockPruneDatasets,
134 mockQueryDatasets_init,
135 mockQueryDatasets_getDatasets,
136 mockQueryDatasets_getTables,
137 cliArgs,
138 exMsgs,
139 exPruneDatasetsCallArgs,
140 exGetTablesCalled,
141 exQueryDatasetsCallArgs,
142 invokeInput=None,
143 exPruneDatasetsExitCode=0,
144 ):
145 """Execute the test.
147 Makes a temporary repo, invokes ``prune-datasets``. Verifies expected
148 output, exit codes, and mock calls.
150 Parameters
151 ----------
152 mockPruneDatasets : `MagicMock`
153 The MagicMock for the ``Butler.pruneDatasets`` function.
154 mockQueryDatasets_init : `MagicMock`
155 The MagicMock for the ``QueryDatasets.__init__`` function.
156 mockQueryDatasets_getDatasets : `MagicMock`
157 The MagicMock for the ``QueryDatasets.getDatasets`` function.
158 mockQueryDatasets_getTables : `MagicMock`
159 The MagicMock for the ``QueryDatasets.getTables`` function.
160 cliArgs : `list` [`str`]
161 The arguments to pass to the command line. Do not include the
162 subcommand name or the repo.
163 exMsgs : `list` [`str`] or None
164 A list of text fragments that should appear in the text output
165 after calling the CLI command, or None if no output should be
166 produced.
167 exPruneDatasetsCallArgs : `dict` [`str`, `Any`]
168 The arguments that ``Butler.pruneDatasets`` should have been called
169 with, or None if that function should not have been called.
170 exGetTablesCalled : bool
171 `True` if ``QueryDatasets.getTables`` should have been called, else
172 `False`.
173 exQueryDatasetsCallArgs : `dict` [`str`, `Any`]
174 The arguments that ``QueryDatasets.__init__`` should have bene
175 called with, or `None` if the function should not have been called.
176 invokeInput : `str`, optional.
177 As string to pass to the ``CliRunner.invoke`` `input` argument. By
178 default None.
179 exPruneDatasetsExitCode : `int`
180 The expected exit code returned from invoking ``prune-datasets``.
181 """
182 runner = LogCliRunner()
183 with runner.isolated_filesystem():
184 # Make a repo so a butler can be created
185 result = runner.invoke(butlerCli, ["create", self.repo])
186 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
188 # Run the prune-datasets CLI command, this will call all of our
189 # mocks:
190 cliArgs = ["prune-datasets", self.repo] + cliArgs
191 result = runner.invoke(butlerCli, cliArgs, input=invokeInput)
192 self.assertEqual(result.exit_code, exPruneDatasetsExitCode, clickResultMsg(result))
194 # Verify the Butler.pruneDatasets was called exactly once with
195 # expected arguments. The datasets argument is the value returned
196 # by QueryDatasets, which we've mocked with side effect
197 # ``getDatasets()``.
198 if exPruneDatasetsCallArgs:
199 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs)
200 else:
201 mockPruneDatasets.assert_not_called()
203 # Less critical, but do a quick verification that the QueryDataset
204 # member function mocks were called, in this case we expect one
205 # time each.
206 if exQueryDatasetsCallArgs:
207 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs)
208 else:
209 mockQueryDatasets_init.assert_not_called()
210 # If Butler.pruneDatasets was not called, then
211 # QueryDatasets.getDatasets also does not get called.
212 if exPruneDatasetsCallArgs:
213 mockQueryDatasets_getDatasets.assert_called_once()
214 else:
215 mockQueryDatasets_getDatasets.assert_not_called()
216 if exGetTablesCalled:
217 mockQueryDatasets_getTables.assert_called_once()
218 else:
219 mockQueryDatasets_getTables.assert_not_called()
221 if exMsgs is None:
222 self.assertEqual("", result.output)
223 else:
224 for expectedMsg in exMsgs:
225 self.assertIn(expectedMsg, result.output)
227 def test_defaults_doContinue(self):
228 """Test running with the default values.
230 Verify that with the default flags that the subcommand says what it
231 will do, prompts for input, and says that it's done.
232 """
233 self.run_test(
234 cliArgs=["myCollection", "--unstore"],
235 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getRefs(), unstore=True),
236 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
237 exGetTablesCalled=True,
238 exMsgs=(
239 pruneDatasets_willRemoveMsg,
240 pruneDatasets_askContinueMsg,
241 astropyTablesToStr(getTables()),
242 pruneDatasets_didRemoveAforementioned,
243 ),
244 invokeInput="yes",
245 )
247 def test_defaults_doNotContinue(self):
248 """Test running with the default values but not continuing.
250 Verify that with the default flags that the subcommand says what it
251 will do, prompts for input, and aborts when told not to continue.
252 """
253 self.run_test(
254 cliArgs=["myCollection", "--unstore"],
255 exPruneDatasetsCallArgs=None,
256 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
257 exGetTablesCalled=True,
258 exMsgs=(
259 pruneDatasets_willRemoveMsg,
260 pruneDatasets_askContinueMsg,
261 pruneDatasets_didNotRemoveAforementioned,
262 ),
263 invokeInput="no",
264 )
266 def test_dryRun_unstore(self):
267 """Test the --dry-run flag with --unstore.
269 Verify that with the dry-run flag the subcommand says what it would
270 remove, but does not remove the datasets.
271 """
272 self.run_test(
273 cliArgs=["myCollection", "--dry-run", "--unstore"],
274 exPruneDatasetsCallArgs=None,
275 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
276 exGetTablesCalled=True,
277 exMsgs=(pruneDatasets_wouldRemoveMsg, astropyTablesToStr(getTables())),
278 )
280 def test_dryRun_disassociate(self):
281 """Test the --dry-run flag with --disassociate.
283 Verify that with the dry-run flag the subcommand says what it would
284 remove, but does not remove the datasets.
285 """
286 collection = "myCollection"
287 self.run_test(
288 cliArgs=[collection, "--dry-run", "--disassociate", "tag1"],
289 exPruneDatasetsCallArgs=None,
290 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=(collection,)),
291 exGetTablesCalled=True,
292 exMsgs=(
293 pruneDatasets_wouldDisassociateMsg.format(collections=(collection,)),
294 astropyTablesToStr(getTables()),
295 ),
296 )
298 def test_dryRun_unstoreAndDisassociate(self):
299 """Test the --dry-run flag with --unstore and --disassociate.
301 Verify that with the dry-run flag the subcommand says what it would
302 remove, but does not remove the datasets.
303 """
304 collection = "myCollection"
305 self.run_test(
306 cliArgs=[collection, "--dry-run", "--unstore", "--disassociate", "tag1"],
307 exPruneDatasetsCallArgs=None,
308 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=(collection,)),
309 exGetTablesCalled=True,
310 exMsgs=(
311 pruneDatasets_wouldDisassociateAndRemoveMsg.format(collections=(collection,)),
312 astropyTablesToStr(getTables()),
313 ),
314 )
316 def test_noConfirm(self):
317 """Test the --no-confirm flag.
319 Verify that with the no-confirm flag the subcommand does not ask for
320 a confirmation, prints the did remove message and the tables that were
321 passed for removal.
322 """
323 self.run_test(
324 cliArgs=["myCollection", "--no-confirm", "--unstore"],
325 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getRefs(), unstore=True),
326 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
327 exGetTablesCalled=True,
328 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
329 )
331 def test_quiet(self):
332 """Test the --quiet flag.
334 Verify that with the quiet flag and the no-confirm flags set that no
335 output is produced by the subcommand.
336 """
337 self.run_test(
338 cliArgs=["myCollection", "--quiet", "--unstore"],
339 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getRefs(), unstore=True),
340 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
341 exGetTablesCalled=True,
342 exMsgs=None,
343 )
345 def test_quietWithDryRun(self):
346 """Test for an error using the --quiet flag with --dry-run."""
347 self.run_test(
348 cliArgs=["--quiet", "--dry-run", "--unstore"],
349 exPruneDatasetsCallArgs=None,
350 exQueryDatasetsCallArgs=None,
351 exGetTablesCalled=False,
352 exMsgs=(pruneDatasets_errQuietWithDryRun,),
353 exPruneDatasetsExitCode=1,
354 )
356 def test_noCollections(self):
357 """Test for an error if no collections are indicated."""
358 self.run_test(
359 cliArgs=["--find-all", "--unstore"],
360 exPruneDatasetsCallArgs=None,
361 exQueryDatasetsCallArgs=None,
362 exGetTablesCalled=False,
363 exMsgs=(pruneDatasets_errNoCollectionRestriction,),
364 exPruneDatasetsExitCode=1,
365 )
367 def test_noDatasets(self):
368 """Test for expected outputs when no datasets are found."""
369 global doFindTables
370 reset = doFindTables
371 try:
372 doFindTables = False
373 self.run_test(
374 cliArgs=["myCollection", "--unstore"],
375 exPruneDatasetsCallArgs=None,
376 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("myCollection",)),
377 exGetTablesCalled=True,
378 exMsgs=(pruneDatasets_noDatasetsFound,),
379 )
380 finally:
381 doFindTables = reset
383 def test_purgeWithDisassociate(self):
384 """Verify there is an error when --purge and --disassociate are both
385 passed in.
386 """
387 self.run_test(
388 cliArgs=["--purge", "run", "--disassociate", "tag1", "tag2"],
389 exPruneDatasetsCallArgs=None,
390 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
391 exGetTablesCalled=False, # ...or this.
392 exMsgs=(pruneDatasets_errPurgeAndDisassociate,),
393 exPruneDatasetsExitCode=1,
394 )
396 def test_purgeNoOp(self):
397 """Verify there is an error when none of --purge, --unstore, or
398 --disassociate are passed.
399 """
400 self.run_test(
401 cliArgs=[],
402 exPruneDatasetsCallArgs=None,
403 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
404 exGetTablesCalled=False, # ...or this.
405 exMsgs=(pruneDatasets_errNoOp,),
406 exPruneDatasetsExitCode=1,
407 )
409 @patch.object(
410 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
411 "get_info",
412 side_effect=lambda x: CollectionInfo(name="run", type=CollectionType.RUN),
413 )
414 def test_purgeImpliedArgs(self, mockGetCollectionType):
415 """Verify the arguments implied by --purge.
417 --purge <run> implies the following arguments to butler.pruneDatasets:
418 purge=True, disassociate=True, unstore=True
419 And for QueryDatasets, if COLLECTIONS is not passed then <run> gets
420 used as the value of COLLECTIONS (and when there is a COLLECTIONS
421 value then find_first gets set to True)
422 """
423 self.run_test(
424 cliArgs=["--purge", "run"],
425 invokeInput="yes",
426 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
427 purge=True, refs=getRefs(), disassociate=True, unstore=True
428 ),
429 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("run",), find_first=True),
430 exGetTablesCalled=True,
431 exMsgs=(
432 pruneDatasets_willRemoveMsg,
433 pruneDatasets_askContinueMsg,
434 astropyTablesToStr(getTables()),
435 pruneDatasets_didRemoveAforementioned,
436 ),
437 )
439 @patch.object(
440 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
441 "get_info",
442 side_effect=lambda x: CollectionInfo(name="run", type=CollectionType.RUN),
443 )
444 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType):
445 """Verify the arguments implied by --purge, with a COLLECTIONS."""
446 self.run_test(
447 cliArgs=["myCollection", "--purge", "run"],
448 invokeInput="yes",
449 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
450 purge=True, disassociate=True, unstore=True, refs=getRefs()
451 ),
452 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
453 collections=("myCollection",), find_first=True
454 ),
455 exGetTablesCalled=True,
456 exMsgs=(
457 pruneDatasets_willRemoveMsg,
458 pruneDatasets_askContinueMsg,
459 astropyTablesToStr(getTables()),
460 pruneDatasets_didRemoveAforementioned,
461 ),
462 )
464 @patch.object(
465 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
466 "get_info",
467 side_effect=lambda x: CollectionInfo(name="myTaggedCollection", type=CollectionType.TAGGED),
468 )
469 def test_purgeOnNonRunCollection(self, mockGetCollectionType):
470 """Verify calling run on a non-run collection fails with expected
471 error message.
472 """
473 collectionName = "myTaggedCollection"
474 self.run_test(
475 cliArgs=["--purge", collectionName],
476 invokeInput="yes",
477 exPruneDatasetsCallArgs=None,
478 exQueryDatasetsCallArgs=None,
479 exGetTablesCalled=False,
480 exMsgs=(pruneDatasets_errPruneOnNotRun.format(collection=collectionName),),
481 exPruneDatasetsExitCode=1,
482 )
484 def test_disassociateImpliedArgs(self):
485 """Verify the arguments implied by --disassociate.
487 --disassociate <tags> implies the following arguments to
488 butler.pruneDatasets:
489 disassociate=True, tags=<tags>
490 and if COLLECTIONS is not passed then <tags> gets used as the value
491 of COLLECTIONS.
493 Use the --no-confirm flag instead of invokeInput="yes", and check for
494 the associated output.
495 """
496 self.run_test(
497 cliArgs=["--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
498 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
499 tags=("tag1", "tag2"), disassociate=True, refs=getRefs()
500 ),
501 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(collections=("tag1", "tag2"), find_first=True),
502 exGetTablesCalled=True,
503 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
504 )
506 def test_disassociateImpliedArgsWithCollections(self):
507 """Verify the arguments implied by --disassociate, with a --collection
508 flag.
509 """
510 self.run_test(
511 cliArgs=["myCollection", "--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
512 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
513 tags=("tag1", "tag2"), disassociate=True, refs=getRefs()
514 ),
515 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
516 collections=("myCollection",), find_first=True
517 ),
518 exGetTablesCalled=True,
519 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
520 )
523if __name__ == "__main__":
524 unittest.main()