Coverage for tests / test_cliCmdRemoveCollections.py: 19%
100 statements
« prev ^ index » next coverage.py v7.13.5, created at 2026-04-26 08:49 +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 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-collections subcommand."""
30import os
31import unittest
32from collections.abc import Sequence
34from astropy.table import Table
35from numpy import array
37from lsst.daf.butler import Butler, CollectionType
38from lsst.daf.butler.cli.butler import cli as butlerCli
39from lsst.daf.butler.cli.cmd._remove_collections import (
40 abortedMsg,
41 canNotRemoveFoundRuns,
42 didNotRemoveFoundRuns,
43 noNonRunCollectionsMsg,
44 removedCollectionsMsg,
45 willRemoveCollectionChainsMsg,
46 willRemoveCollectionMsg,
47)
48from lsst.daf.butler.cli.utils import LogCliRunner, clickResultMsg
49from lsst.daf.butler.script.removeCollections import removeCollections
50from lsst.daf.butler.tests.utils import (
51 ButlerTestHelper,
52 MetricTestRepo,
53 makeTestTempDir,
54 readTable,
55 removeTestTempDir,
56)
58TESTDIR = os.path.abspath(os.path.dirname(__file__))
60QueryCollectionsRow = tuple[str, str] | tuple[str, str, str]
61RemoveCollectionRow = tuple[str, str]
64class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper):
65 """Test executing remove collection."""
67 def setUp(self):
68 self.runner = LogCliRunner()
70 self.root = makeTestTempDir(TESTDIR)
71 self.testRepo = MetricTestRepo(
72 self.root, configFile=os.path.join(TESTDIR, "config/basic/butler.yaml")
73 )
74 self.enterContext(self.testRepo.butler)
76 def tearDown(self):
77 removeTestTempDir(self.root)
79 def _verify_remove(
80 self,
81 collection: str,
82 before_rows: Sequence[QueryCollectionsRow],
83 remove_rows: Sequence[RemoveCollectionRow],
84 after_rows: Sequence[QueryCollectionsRow],
85 ):
86 """Remove collections, with verification that expected collections are
87 present before removing, that the command reports expected collections
88 to be removed, and that expected collections are present after removal.
90 Parameters
91 ----------
92 collection : `str`
93 The name of the collection, or glob pattern for collections, to
94 remove.
95 before_rows : `~collections.abc.Sequence` [ `QueryCollectionsRow` ]
96 The rows that should be in the table returned by query-collections
97 before removing the collection.
98 remove_rows : `~collections.abc.Sequence` [ `RemoveCollectionRow` ]
99 The rows that should be in the "will remove" table while removing
100 collections.
101 after_rows : `~collections.abc.Sequence` [ `QueryCollectionsRow` ]
102 The rows that should be in the table returned by query-collections
103 after removing the collection.
104 """
106 def _query_collection_column_names(rows):
107 # If there is a chained collection in the table then there is a
108 # definition column, otherwise there is only the name and type
109 # columns.
110 if len(rows[0]) == 2:
111 return ("Name", "Type")
112 elif len(rows[0]) == 3:
113 return ("Name", "Type", "Children")
114 else:
115 raise RuntimeError(f"Unhandled column count: {len(rows[0])}")
117 result = self.runner.invoke(butlerCli, ["query-collections", self.root, "--chains", "TABLE"])
118 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
119 expected = Table(array(before_rows), names=_query_collection_column_names(before_rows))
120 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
122 removal = removeCollections(repo=self.root, collection=collection, remove_from_parents=False)
123 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
124 expected = Table(array(remove_rows), names=("Collection", "Collection Type"))
125 self.assertAstropyTablesEqual(removal.removeCollectionsTable, expected)
126 removal.onConfirmation()
128 result = self.runner.invoke(butlerCli, ["query-collections", self.root])
129 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
130 expected = Table(array(after_rows), names=_query_collection_column_names(after_rows))
131 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
133 def testRemoveScript(self):
134 """Test removing collections.
136 Combining several tests into one case allows us to reuse the test repo,
137 which saves execution time.
138 """
139 # Test wildcard with chained collections:
141 # Add a couple chained collections
142 for parent, child in (
143 ("chained-run-1", "ingest/run"),
144 ("chained-run-2", "ingest/run"),
145 ):
146 result = self.runner.invoke(
147 butlerCli,
148 ["collection-chain", self.root, parent, child],
149 )
150 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
152 self._verify_remove(
153 collection="chained-run-*",
154 before_rows=(
155 ("chained-run-1", "CHAINED", "ingest/run"),
156 ("chained-run-2", "CHAINED", "ingest/run"),
157 ("ingest", "TAGGED", ""),
158 ("ingest/run", "RUN", ""),
159 ),
160 remove_rows=(
161 ("chained-run-1", "CHAINED"),
162 ("chained-run-2", "CHAINED"),
163 ),
164 after_rows=(
165 ("ingest", "TAGGED"),
166 ("ingest/run", "RUN"),
167 ),
168 )
170 # Test a single tagged collection:
172 self._verify_remove(
173 collection="ingest",
174 before_rows=(
175 ("ingest", "TAGGED"),
176 ("ingest/run", "RUN"),
177 ),
178 remove_rows=(("ingest", "TAGGED"),),
179 after_rows=(("ingest/run", "RUN"),),
180 )
182 def testRemoveCmd(self):
183 """Test remove command outputs."""
184 # Test expected output with a non-existent collection:
186 result = self.runner.invoke(butlerCli, ["remove-collections", self.root, "fake_collection"])
187 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
188 self.assertIn(noNonRunCollectionsMsg, result.stdout)
190 # Add a couple chained collections
191 for parent, child in (
192 ("chained-run-1", "ingest/run"),
193 ("chained-run-2", "ingest/run"),
194 ):
195 result = self.runner.invoke(
196 butlerCli,
197 ["collection-chain", self.root, parent, child],
198 )
199 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
201 # Test aborting a removal
203 result = self.runner.invoke(
204 butlerCli,
205 ["remove-collections", self.root, "chained-run-1"],
206 input="no",
207 )
208 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
209 self.assertIn(abortedMsg, result.stdout)
211 # Remove with --no-confirm, it's expected to run silently.
213 result = self.runner.invoke(
214 butlerCli, ["remove-collections", self.root, "chained-run-1", "--no-confirm"]
215 )
216 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
217 self.assertIn(removedCollectionsMsg, result.stdout)
218 self.assertIn("chained-run-1", result.stdout)
220 # verify chained-run-1 was removed:
222 butler = Butler.from_config(self.root)
223 self.enterContext(butler)
224 collections = butler.registry.queryCollections(
225 collectionTypes=frozenset(
226 (
227 CollectionType.RUN,
228 CollectionType.TAGGED,
229 CollectionType.CHAINED,
230 CollectionType.CALIBRATION,
231 )
232 ),
233 )
234 self.assertCountEqual(["ingest/run", "ingest", "chained-run-2"], collections)
236 # verify chained-run-2 can be removed with prompting and expected CLI
237 # output
239 result = self.runner.invoke(
240 butlerCli,
241 ["remove-collections", self.root, "chained-run-2"],
242 input="yes",
243 )
244 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
245 self.assertIn(willRemoveCollectionMsg, result.stdout)
246 self.assertIn("chained-run-2 CHAINED", result.stdout)
248 # try to remove a run table, check for the "can not remove run" message
250 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child])
251 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
252 result = self.runner.invoke(
253 # removes run-chain (chained collection), but can not remove the
254 # run collection, and emits a message that says so.
255 butlerCli,
256 ["remove-collections", self.root, "*run*"],
257 input="yes",
258 )
259 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
260 self.assertIn(canNotRemoveFoundRuns, result.stdout)
261 self.assertIn("ingest/run", result.stdout)
263 # try to remove a run table with --no-confirm, check for the "did not
264 # remove run" message
266 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child])
267 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
268 result = self.runner.invoke(
269 # removes run-chain (chained collection), but can not remove the
270 # run collection, and emits a message that says so.
271 butlerCli,
272 ["remove-collections", self.root, "*run*", "--no-confirm"],
273 )
274 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
275 self.assertIn(didNotRemoveFoundRuns, result.stdout)
276 self.assertIn("ingest/run", result.stdout)
278 def testRemoveFromParents(self) -> None:
279 butler = Butler(self.root, writeable=True)
280 self.enterContext(butler)
281 butler.collections.register("tag1", CollectionType.TAGGED)
282 butler.collections.register("tag2", CollectionType.TAGGED)
283 butler.collections.register("chain1", CollectionType.CHAINED)
284 butler.collections.register("chain2", CollectionType.CHAINED)
285 butler.collections.register("chain3", CollectionType.CHAINED)
286 butler.collections.redefine_chain("chain1", ["tag1", "tag2", "chain3"])
287 butler.collections.redefine_chain("chain2", ["tag1"])
289 # Make sure the printed output is correct.
290 removal = removeCollections(repo=self.root, collection="tag*", remove_from_parents=True)
291 table = [tuple(row) for row in removal.removeChainsTable]
292 self.assertEqual(table, [("tag1", "chain1"), ("", "chain2"), ("tag2", "chain1")])
293 result = self.runner.invoke(
294 butlerCli,
295 ["remove-collections", "--remove-from-parents", self.root, "tag*"],
296 input="yes",
297 )
298 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
299 self.assertIn(willRemoveCollectionChainsMsg, result.stdout)
300 # Make sure the collections are deleted as expected.
301 self.assertEqual(
302 sorted(butler.collections.query("*")), ["chain1", "chain2", "chain3", "ingest", "ingest/run"]
303 )
304 self.assertEqual(butler.collections.get_info("chain1").children, ("chain3",))
307if __name__ == "__main__":
308 unittest.main()