Coverage for tests/test_cliCmdPruneDatasets.py: 39%
106 statements
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-17 09:33 +0000
« prev ^ index » next coverage.py v7.2.5, created at 2023-05-17 09:33 +0000
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_errNoOp,
41 pruneDatasets_errPruneOnNotRun,
42 pruneDatasets_errPurgeAndDisassociate,
43 pruneDatasets_errQuietWithDryRun,
44 pruneDatasets_noDatasetsFound,
45 pruneDatasets_willRemoveMsg,
46 pruneDatasets_wouldDisassociateAndRemoveMsg,
47 pruneDatasets_wouldDisassociateMsg,
48 pruneDatasets_wouldRemoveMsg,
49)
50from lsst.daf.butler.cli.utils import LogCliRunner, astropyTablesToStr, clickResultMsg
51from lsst.daf.butler.registry import CollectionType
52from lsst.daf.butler.script import QueryDatasets
54doFindTables = True
57def getTables():
58 if doFindTables:
59 return (Table(((1, 2, 3),), names=("foo",)),)
60 return tuple()
63def getDatasets():
64 return "datasets"
67def makeQueryDatasets(*args, **kwargs):
68 return QueryDatasets(*args, **kwargs)
71class PruneDatasetsTestCase(unittest.TestCase):
72 """Tests the ``prune_datasets`` "command" function (in
73 ``cli/cmd/commands.py``) and the ``pruneDatasets`` "script" function (in
74 ``scripts/_pruneDatasets.py``).
76 ``Butler.pruneDatasets`` and a few other functions that get called before
77 it are mocked, and tests check for expected arguments to those mocks.
78 """
80 def setUp(self):
81 self.repo = "here"
83 @staticmethod
84 def makeQueryDatasetsArgs(*, repo, **kwargs):
85 expectedArgs = dict(
86 repo=repo, collections=(), where="", find_first=True, show_uri=False, glob=tuple()
87 )
88 expectedArgs.update(kwargs)
89 return expectedArgs
91 @staticmethod
92 def makePruneDatasetsArgs(**kwargs):
93 expectedArgs = dict(refs=tuple(), disassociate=False, tags=(), purge=False, unstore=False)
94 expectedArgs.update(kwargs)
95 return expectedArgs
97 # Mock the QueryDatasets.getTables function to return a set of Astropy
98 # tables, similar to what would be returned by a call to
99 # QueryDatasets.getTables on a repo with real data.
100 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getTables", side_effect=getTables)
101 # Mock the QueryDatasets.getDatasets function. Normally it would return a
102 # list of queries.DatasetQueryResults, but all we need to do is verify that
103 # the output of this function is passed into our pruneDatasets magicMock,
104 # so we can return something arbitrary that we can test is equal."""
105 @patch.object(lsst.daf.butler.script._pruneDatasets.QueryDatasets, "getDatasets", side_effect=getDatasets)
106 # Mock the actual QueryDatasets class, so we can inspect calls to its init
107 # function. Note that the side_effect returns an instance of QueryDatasets,
108 # so this mock records and then is a pass-through.
109 @patch.object(lsst.daf.butler.script._pruneDatasets, "QueryDatasets", side_effect=makeQueryDatasets)
110 # Mock the pruneDatasets butler command so we can test for expected calls
111 # to it, without dealing with setting up a full repo with data for it.
112 @patch.object(Butler, "pruneDatasets")
113 def run_test(
114 self,
115 mockPruneDatasets,
116 mockQueryDatasets_init,
117 mockQueryDatasets_getDatasets,
118 mockQueryDatasets_getTables,
119 cliArgs,
120 exMsgs,
121 exPruneDatasetsCallArgs,
122 exGetTablesCalled,
123 exQueryDatasetsCallArgs,
124 invokeInput=None,
125 exPruneDatasetsExitCode=0,
126 ):
127 """Execute the test.
129 Makes a temporary repo, invokes ``prune-datasets``. Verifies expected
130 output, exit codes, and mock calls.
132 Parameters
133 ----------
134 mockPruneDatasets : `MagicMock`
135 The MagicMock for the ``Butler.pruneDatasets`` function.
136 mockQueryDatasets_init : `MagicMock`
137 The MagicMock for the ``QueryDatasets.__init__`` function.
138 mockQueryDatasets_getDatasets : `MagicMock`
139 The MagicMock for the ``QueryDatasets.getDatasets`` function.
140 mockQueryDatasets_getTables : `MagicMock`
141 The MagicMock for the ``QueryDatasets.getTables`` function.
142 cliArgs : `list` [`str`]
143 The arguments to pass to the command line. Do not include the
144 subcommand name or the repo.
145 exMsgs : `list` [`str`] or None
146 A list of text fragments that should appear in the text output
147 after calling the CLI command, or None if no output should be
148 produced.
149 exPruneDatasetsCallArgs : `dict` [`str`, `Any`]
150 The arguments that ``Butler.pruneDatasets`` should have been called
151 with, or None if that function should not have been called.
152 exGetTablesCalled : bool
153 `True` if ``QueryDatasets.getTables`` should have been called, else
154 `False`.
155 exQueryDatasetsCallArgs : `dict` [`str`, `Any`]
156 The arguments that ``QueryDatasets.__init__`` should have bene
157 called with, or `None` if the function should not have been called.
158 invokeInput : `str`, optional.
159 As string to pass to the ``CliRunner.invoke`` `input` argument. By
160 default None.
161 exPruneDatasetsExitCode : `int`
162 The expected exit code returned from invoking ``prune-datasets``.
163 """
164 runner = LogCliRunner()
165 with runner.isolated_filesystem():
166 # Make a repo so a butler can be created
167 result = runner.invoke(butlerCli, ["create", self.repo])
168 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
170 # Run the prune-datasets CLI command, this will call all of our
171 # mocks:
172 cliArgs = ["prune-datasets", self.repo] + cliArgs
173 result = runner.invoke(butlerCli, cliArgs, input=invokeInput)
174 self.assertEqual(result.exit_code, exPruneDatasetsExitCode, clickResultMsg(result))
176 # Verify the Butler.pruneDatasets was called exactly once with
177 # expected arguments. The datasets argument is the value returned
178 # by QueryDatasets, which we've mocked with side effect
179 # ``getDatasets()``.
180 if exPruneDatasetsCallArgs:
181 mockPruneDatasets.assert_called_once_with(**exPruneDatasetsCallArgs)
182 else:
183 mockPruneDatasets.assert_not_called()
185 # Less critical, but do a quick verification that the QueryDataset
186 # member function mocks were called, in this case we expect one
187 # time each.
188 if exQueryDatasetsCallArgs:
189 mockQueryDatasets_init.assert_called_once_with(**exQueryDatasetsCallArgs)
190 else:
191 mockQueryDatasets_init.assert_not_called()
192 # If Butler.pruneDatasets was not called, then
193 # QueryDatasets.getDatasets also does not get called.
194 if exPruneDatasetsCallArgs:
195 mockQueryDatasets_getDatasets.assert_called_once()
196 else:
197 mockQueryDatasets_getDatasets.assert_not_called()
198 if exGetTablesCalled:
199 mockQueryDatasets_getTables.assert_called_once()
200 else:
201 mockQueryDatasets_getTables.assert_not_called()
203 if exMsgs is None:
204 self.assertEqual("", result.output)
205 else:
206 for expectedMsg in exMsgs:
207 self.assertIn(expectedMsg, result.output)
209 def test_defaults_doContinue(self):
210 """Test running with the default values.
212 Verify that with the default flags that the subcommand says what it
213 will do, prompts for input, and says that it's done."""
214 self.run_test(
215 cliArgs=["myCollection", "--unstore"],
216 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True),
217 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)),
218 exGetTablesCalled=True,
219 exMsgs=(
220 pruneDatasets_willRemoveMsg,
221 pruneDatasets_askContinueMsg,
222 astropyTablesToStr(getTables()),
223 pruneDatasets_didRemoveAforementioned,
224 ),
225 invokeInput="yes",
226 )
228 def test_defaults_doNotContinue(self):
229 """Test running with the default values but not continuing.
231 Verify that with the default flags that the subcommand says what it
232 will do, prompts for input, and aborts when told not to continue."""
233 self.run_test(
234 cliArgs=["myCollection", "--unstore"],
235 exPruneDatasetsCallArgs=None,
236 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)),
237 exGetTablesCalled=True,
238 exMsgs=(
239 pruneDatasets_willRemoveMsg,
240 pruneDatasets_askContinueMsg,
241 pruneDatasets_didNotRemoveAforementioned,
242 ),
243 invokeInput="no",
244 )
246 def test_dryRun_unstore(self):
247 """Test the --dry-run flag with --unstore.
249 Verify that with the dry-run flag the subcommand says what it would
250 remove, but does not remove the datasets."""
251 self.run_test(
252 cliArgs=["myCollection", "--dry-run", "--unstore"],
253 exPruneDatasetsCallArgs=None,
254 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)),
255 exGetTablesCalled=True,
256 exMsgs=(pruneDatasets_wouldRemoveMsg, astropyTablesToStr(getTables())),
257 )
259 def test_dryRun_disassociate(self):
260 """Test the --dry-run flag with --disassociate.
262 Verify that with the dry-run flag the subcommand says what it would
263 remove, but does not remove the datasets."""
264 collection = "myCollection"
265 self.run_test(
266 cliArgs=[collection, "--dry-run", "--disassociate", "tag1"],
267 exPruneDatasetsCallArgs=None,
268 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=(collection,)),
269 exGetTablesCalled=True,
270 exMsgs=(
271 pruneDatasets_wouldDisassociateMsg.format(collections=(collection,)),
272 astropyTablesToStr(getTables()),
273 ),
274 )
276 def test_dryRun_unstoreAndDisassociate(self):
277 """Test the --dry-run flag with --unstore and --disassociate.
279 Verify that with the dry-run flag the subcommand says what it would
280 remove, but does not remove the datasets."""
281 collection = "myCollection"
282 self.run_test(
283 cliArgs=[collection, "--dry-run", "--unstore", "--disassociate", "tag1"],
284 exPruneDatasetsCallArgs=None,
285 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=(collection,)),
286 exGetTablesCalled=True,
287 exMsgs=(
288 pruneDatasets_wouldDisassociateAndRemoveMsg.format(collections=(collection,)),
289 astropyTablesToStr(getTables()),
290 ),
291 )
293 def test_noConfirm(self):
294 """Test the --no-confirm flag.
296 Verify that with the no-confirm flag the subcommand does not ask for
297 a confirmation, prints the did remove message and the tables that were
298 passed for removal."""
299 self.run_test(
300 cliArgs=["myCollection", "--no-confirm", "--unstore"],
301 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True),
302 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)),
303 exGetTablesCalled=True,
304 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
305 )
307 def test_quiet(self):
308 """Test the --quiet flag.
310 Verify that with the quiet flag and the no-confirm flags set that no
311 output is produced by the subcommand."""
312 self.run_test(
313 cliArgs=["myCollection", "--quiet", "--unstore"],
314 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(refs=getDatasets(), unstore=True),
315 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(repo=self.repo, collections=("myCollection",)),
316 exGetTablesCalled=True,
317 exMsgs=None,
318 )
320 def test_quietWithDryRun(self):
321 """Test for an error using the --quiet flag with --dry-run."""
322 self.run_test(
323 cliArgs=["--quiet", "--dry-run", "--unstore"],
324 exPruneDatasetsCallArgs=None,
325 exQueryDatasetsCallArgs=None,
326 exGetTablesCalled=False,
327 exMsgs=(pruneDatasets_errQuietWithDryRun,),
328 exPruneDatasetsExitCode=1,
329 )
331 def test_noCollections(self):
332 """Test for an error if no collections are indicated."""
333 self.run_test(
334 cliArgs=["--find-all", "--unstore"],
335 exPruneDatasetsCallArgs=None,
336 exQueryDatasetsCallArgs=None,
337 exGetTablesCalled=False,
338 exMsgs=(pruneDatasets_errNoCollectionRestriction,),
339 exPruneDatasetsExitCode=1,
340 )
342 def test_noDatasets(self):
343 """Test for expected outputs when no datasets are found."""
344 global doFindTables
345 reset = doFindTables
346 try:
347 doFindTables = False
348 self.run_test(
349 cliArgs=["myCollection", "--unstore"],
350 exPruneDatasetsCallArgs=None,
351 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
352 repo=self.repo, collections=("myCollection",)
353 ),
354 exGetTablesCalled=True,
355 exMsgs=(pruneDatasets_noDatasetsFound,),
356 )
357 finally:
358 doFindTables = reset
360 def test_purgeWithDisassociate(self):
361 """Verify there is an error when --purge and --disassociate are both
362 passed in."""
363 self.run_test(
364 cliArgs=["--purge", "run", "--disassociate", "tag1", "tag2"],
365 exPruneDatasetsCallArgs=None,
366 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
367 exGetTablesCalled=False, # ...or this.
368 exMsgs=(pruneDatasets_errPurgeAndDisassociate,),
369 exPruneDatasetsExitCode=1,
370 )
372 def test_purgeNoOp(self):
373 """Verify there is an error when none of --purge, --unstore, or
374 --disassociate are passed."""
375 self.run_test(
376 cliArgs=[],
377 exPruneDatasetsCallArgs=None,
378 exQueryDatasetsCallArgs=None, # should not make it far enough to call this.
379 exGetTablesCalled=False, # ...or this.
380 exMsgs=(pruneDatasets_errNoOp,),
381 exPruneDatasetsExitCode=1,
382 )
384 @patch.object( 384 ↛ exitline 384 didn't jump to the function exit
385 lsst.daf.butler.registries.sql.SqlRegistry,
386 "getCollectionType",
387 side_effect=lambda x: CollectionType.RUN,
388 )
389 def test_purgeImpliedArgs(self, mockGetCollectionType):
390 """Verify the arguments implied by --purge.
392 --purge <run> implies the following arguments to butler.pruneDatasets:
393 purge=True, disassociate=True, unstore=True
394 And for QueryDatasets, if COLLECTIONS is not passed then <run> gets
395 used as the value of COLLECTIONS (and when there is a COLLECTIONS
396 value then find_first gets set to True)
397 """
398 self.run_test(
399 cliArgs=["--purge", "run"],
400 invokeInput="yes",
401 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
402 purge=True, refs=getDatasets(), disassociate=True, unstore=True
403 ),
404 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
405 repo=self.repo, collections=("run",), find_first=True
406 ),
407 exGetTablesCalled=True,
408 exMsgs=(
409 pruneDatasets_willRemoveMsg,
410 pruneDatasets_askContinueMsg,
411 astropyTablesToStr(getTables()),
412 pruneDatasets_didRemoveAforementioned,
413 ),
414 )
416 @patch.object( 416 ↛ exitline 416 didn't jump to the function exit
417 lsst.daf.butler.registries.sql.SqlRegistry,
418 "getCollectionType",
419 side_effect=lambda x: CollectionType.RUN,
420 )
421 def test_purgeImpliedArgsWithCollections(self, mockGetCollectionType):
422 """Verify the arguments implied by --purge, with a COLLECTIONS."""
423 self.run_test(
424 cliArgs=["myCollection", "--purge", "run"],
425 invokeInput="yes",
426 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
427 purge=True, disassociate=True, unstore=True, refs=getDatasets()
428 ),
429 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
430 repo=self.repo, collections=("myCollection",), find_first=True
431 ),
432 exGetTablesCalled=True,
433 exMsgs=(
434 pruneDatasets_willRemoveMsg,
435 pruneDatasets_askContinueMsg,
436 astropyTablesToStr(getTables()),
437 pruneDatasets_didRemoveAforementioned,
438 ),
439 )
441 @patch.object( 441 ↛ exitline 441 didn't jump to the function exit
442 lsst.daf.butler.registries.sql.SqlRegistry,
443 "getCollectionType",
444 side_effect=lambda x: CollectionType.TAGGED,
445 )
446 def test_purgeOnNonRunCollection(self, mockGetCollectionType):
447 """Verify calling run on a non-run collection fails with expected
448 error message."""
449 collectionName = "myTaggedCollection"
450 self.run_test(
451 cliArgs=["--purge", collectionName],
452 invokeInput="yes",
453 exPruneDatasetsCallArgs=None,
454 exQueryDatasetsCallArgs=None,
455 exGetTablesCalled=False,
456 exMsgs=(pruneDatasets_errPruneOnNotRun.format(collection=collectionName)),
457 exPruneDatasetsExitCode=1,
458 )
460 def test_disassociateImpliedArgs(self):
461 """Verify the arguments implied by --disassociate.
463 --disassociate <tags> implies the following arguments to
464 butler.pruneDatasets:
465 disassociate=True, tags=<tags>
466 and if COLLECTIONS is not passed then <tags> gets used as the value
467 of COLLECTIONS.
469 Use the --no-confirm flag instead of invokeInput="yes", and check for
470 the associated output.
471 """
472 self.run_test(
473 cliArgs=["--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
474 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
475 tags=("tag1", "tag2"), disassociate=True, refs=getDatasets()
476 ),
477 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
478 repo=self.repo, collections=("tag1", "tag2"), find_first=True
479 ),
480 exGetTablesCalled=True,
481 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
482 )
484 def test_disassociateImpliedArgsWithCollections(self):
485 """Verify the arguments implied by --disassociate, with a --collection
486 flag."""
487 self.run_test(
488 cliArgs=["myCollection", "--disassociate", "tag1", "--disassociate", "tag2", "--no-confirm"],
489 exPruneDatasetsCallArgs=self.makePruneDatasetsArgs(
490 tags=("tag1", "tag2"), disassociate=True, refs=getDatasets()
491 ),
492 exQueryDatasetsCallArgs=self.makeQueryDatasetsArgs(
493 repo=self.repo, collections=("myCollection",), find_first=True
494 ),
495 exGetTablesCalled=True,
496 exMsgs=(pruneDatasets_didRemoveMsg, astropyTablesToStr(getTables())),
497 )
500if __name__ == "__main__":
501 unittest.main()