Coverage for tests/test_cliCmdPruneDatasets.py: 38%
110 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-20 01:50 -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 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(*, repo, **kwargs):
103 expectedArgs = dict(
104 repo=repo, 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(repo=self.repo, 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(repo=self.repo, 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(repo=self.repo, 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(repo=self.repo, 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(repo=self.repo, 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(repo=self.repo, 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(repo=self.repo, 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(
377 repo=self.repo, collections=("myCollection",)
378 ),
379 exGetTablesCalled=True,
380 exMsgs=(pruneDatasets_noDatasetsFound,),
381 )
382 finally:
383 doFindTables = reset
385 def test_purgeWithDisassociate(self):
386 """Verify there is an error when --purge and --disassociate are both
387 passed in.
388 """
389 self.run_test(
390 cliArgs=["--purge", "run", "--disassociate", "tag1", "tag2"],
391 exPruneDatasetsCallArgs=None,
392 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
393 exGetTablesCalled=False, # ...or this.
394 exMsgs=(pruneDatasets_errPurgeAndDisassociate,),
395 exPruneDatasetsExitCode=1,
396 )
398 def test_purgeNoOp(self):
399 """Verify there is an error when none of --purge, --unstore, or
400 --disassociate are passed.
401 """
402 self.run_test(
403 cliArgs=[],
404 exPruneDatasetsCallArgs=None,
405 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
406 exGetTablesCalled=False, # ...or this.
407 exMsgs=(pruneDatasets_errNoOp,),
408 exPruneDatasetsExitCode=1,
409 )
411 @patch.object(
412 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
413 "get_info",
414 side_effect=lambda x: CollectionInfo(name="run", type=CollectionType.RUN),
415 )
416 def test_purgeImpliedArgs(self, mockGetCollectionType):
417 """Verify the arguments implied by --purge.
419 --purge <run> implies the following arguments to butler.pruneDatasets:
420 purge=True, disassociate=True, unstore=True
421 And for QueryDatasets, if COLLECTIONS is not passed then <run> gets
422 used as the value of COLLECTIONS (and when there is a COLLECTIONS
423 value then find_first gets set to True)
424 """
425 self.run_test(
426 cliArgs=["--purge", "run"],
427 invokeInput="yes",
428 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
429 purge=True, refs=getRefs(), disassociate=True, unstore=True
430 ),
431 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
432 repo=self.repo, collections=("run",), find_first=True
433 ),
434 exGetTablesCalled=True,
435 exMsgs=(
436 pruneDatasets_willRemoveMsg,
437 pruneDatasets_askContinueMsg,
438 astropyTablesToStr(getTables()),
439 pruneDatasets_didRemoveAforementioned,
440 ),
441 )
443 @patch.object(
444 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
445 "get_info",
446 side_effect=lambda x: CollectionInfo(name="run", type=CollectionType.RUN),
447 )
448 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType):
449 """Verify the arguments implied by --purge, with a COLLECTIONS."""
450 self.run_test(
451 cliArgs=["myCollection", "--purge", "run"],
452 invokeInput="yes",
453 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
454 purge=True, disassociate=True, unstore=True, refs=getRefs()
455 ),
456 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
457 repo=self.repo, collections=("myCollection",), find_first=True
458 ),
459 exGetTablesCalled=True,
460 exMsgs=(
461 pruneDatasets_willRemoveMsg,
462 pruneDatasets_askContinueMsg,
463 astropyTablesToStr(getTables()),
464 pruneDatasets_didRemoveAforementioned,
465 ),
466 )
468 @patch.object(
469 lsst.daf.butler.direct_butler._direct_butler_collections.DirectButlerCollections,
470 "get_info",
471 side_effect=lambda x: CollectionInfo(name="myTaggedCollection", type=CollectionType.TAGGED),
472 )
473 def test_purgeOnNonRunCollection(self, mockGetCollectionType):
474 """Verify calling run on a non-run collection fails with expected
475 error message.
476 """
477 collectionName = "myTaggedCollection"
478 self.run_test(
479 cliArgs=["--purge", collectionName],
480 invokeInput="yes",
481 exPruneDatasetsCallArgs=None,
482 exQueryDatasetsCallArgs=None,
483 exGetTablesCalled=False,
484 exMsgs=(pruneDatasets_errPruneOnNotRun.format(collection=collectionName),),
485 exPruneDatasetsExitCode=1,
486 )
488 def test_disassociateImpliedArgs(self):
489 """Verify the arguments implied by --disassociate.
491 --disassociate <tags> implies the following arguments to
492 butler.pruneDatasets:
493 disassociate=True, tags=<tags>
494 and if COLLECTIONS is not passed then <tags> gets used as the value
495 of COLLECTIONS.
497 Use the --no-confirm flag instead of invokeInput="yes", and check for
498 the associated output.
499 """
500 self.run_test(
501 cliArgs=["--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
502 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
503 tags=("tag1", "tag2"), disassociate=True, refs=getRefs()
504 ),
505 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
506 repo=self.repo, collections=("tag1", "tag2"), find_first=True
507 ),
508 exGetTablesCalled=True,
509 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
510 )
512 def test_disassociateImpliedArgsWithCollections(self):
513 """Verify the arguments implied by --disassociate, with a --collection
514 flag.
515 """
516 self.run_test(
517 cliArgs=["myCollection", "--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
518 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
519 tags=("tag1", "tag2"), disassociate=True, refs=getRefs()
520 ),
521 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
522 repo=self.repo, collections=("myCollection",), find_first=True
523 ),
524 exGetTablesCalled=True,
525 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
526 )
529if __name__ == "__main__":
530 unittest.main()