Coverage for tests/test_cliCmdRemoveCollections.py: 24%
81 statements
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:53 -0700
« prev ^ index » next coverage.py v7.4.4, created at 2024-04-05 02:53 -0700
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.
29"""
31import os
32import unittest
33from collections.abc import Sequence
35from astropy.table import Table
36from lsst.daf.butler import Butler, CollectionType
37from lsst.daf.butler.cli.butler import cli as butlerCli
38from lsst.daf.butler.cli.cmd._remove_collections import (
39 abortedMsg,
40 canNotRemoveFoundRuns,
41 didNotRemoveFoundRuns,
42 noNonRunCollectionsMsg,
43 removedCollectionsMsg,
44 willRemoveCollectionMsg,
45)
46from lsst.daf.butler.cli.utils import LogCliRunner, clickResultMsg
47from lsst.daf.butler.script.removeCollections import removeCollections
48from lsst.daf.butler.tests.utils import (
49 ButlerTestHelper,
50 MetricTestRepo,
51 makeTestTempDir,
52 readTable,
53 removeTestTempDir,
54)
55from numpy import array
57TESTDIR = os.path.abspath(os.path.dirname(__file__))
59QueryCollectionsRow = tuple[str, str] | tuple[str, str, str]
60RemoveCollectionRow = tuple[str, str]
63class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper):
64 """Test executing remove collection."""
66 def setUp(self):
67 self.runner = LogCliRunner()
69 self.root = makeTestTempDir(TESTDIR)
70 self.testRepo = MetricTestRepo(
71 self.root, configFile=os.path.join(TESTDIR, "config/basic/butler.yaml")
72 )
74 def tearDown(self):
75 removeTestTempDir(self.root)
77 def _verify_remove(
78 self,
79 collection: str,
80 before_rows: Sequence[QueryCollectionsRow],
81 remove_rows: Sequence[RemoveCollectionRow],
82 after_rows: Sequence[QueryCollectionsRow],
83 ):
84 """Remove collections, with verification that expected collections are
85 present before removing, that the command reports expected collections
86 to be removed, and that expected collections are present after removal.
88 Parameters
89 ----------
90 collection : `str`
91 The name of the collection, or glob pattern for collections, to
92 remove.
93 before_rows : `Sequence` [ `QueryCollectionsRow` ]
94 The rows that should be in the table returned by query-collections
95 before removing the collection.
96 remove_rows : `Sequence` [ `RemoveCollectionRow` ]
97 The rows that should be in the "will remove" table while removing
98 collections.
99 after_rows : `Sequence` [ `QueryCollectionsRow` ]
100 The rows that should be in the table returned by query-collections
101 after removing the collection.
102 """
104 def _query_collection_column_names(rows):
105 # If there is a chained collection in the table then there is a
106 # definition column, otherwise there is only the name and type
107 # columns.
108 if len(rows[0]) == 2:
109 return ("Name", "Type")
110 elif len(rows[0]) == 3:
111 return ("Name", "Type", "Children")
112 else:
113 raise RuntimeError(f"Unhandled column count: {len(rows[0])}")
115 result = self.runner.invoke(butlerCli, ["query-collections", self.root, "--chains", "TABLE"])
116 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
117 expected = Table(array(before_rows), names=_query_collection_column_names(before_rows))
118 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
120 removal = removeCollections(
121 repo=self.root,
122 collection=collection,
123 )
124 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
125 expected = Table(array(remove_rows), names=("Collection", "Collection Type"))
126 self.assertAstropyTablesEqual(removal.removeCollectionsTable, expected)
127 removal.onConfirmation()
129 result = self.runner.invoke(butlerCli, ["query-collections", self.root])
130 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
131 expected = Table(array(after_rows), names=_query_collection_column_names(after_rows))
132 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
134 def testRemoveScript(self):
135 """Test removing collections.
137 Combining several tests into one case allows us to reuse the test repo,
138 which saves execution time.
139 """
140 # Test wildcard with chained collections:
142 # Add a couple chained collections
143 for parent, child in (
144 ("chained-run-1", "ingest/run"),
145 ("chained-run-2", "ingest/run"),
146 ):
147 result = self.runner.invoke(
148 butlerCli,
149 ["collection-chain", self.root, parent, child],
150 )
151 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
153 self._verify_remove(
154 collection="chained-run-*",
155 before_rows=(
156 ("chained-run-1", "CHAINED", "ingest/run"),
157 ("chained-run-2", "CHAINED", "ingest/run"),
158 ("ingest", "TAGGED", ""),
159 ("ingest/run", "RUN", ""),
160 ),
161 remove_rows=(
162 ("chained-run-1", "CHAINED"),
163 ("chained-run-2", "CHAINED"),
164 ),
165 after_rows=(
166 ("ingest", "TAGGED"),
167 ("ingest/run", "RUN"),
168 ),
169 )
171 # Test a single tagged collection:
173 self._verify_remove(
174 collection="ingest",
175 before_rows=(
176 ("ingest", "TAGGED"),
177 ("ingest/run", "RUN"),
178 ),
179 remove_rows=(("ingest", "TAGGED"),),
180 after_rows=(("ingest/run", "RUN"),),
181 )
183 def testRemoveCmd(self):
184 """Test remove command outputs."""
185 # Test expected output with a non-existent collection:
187 result = self.runner.invoke(butlerCli, ["remove-collections", self.root, "fake_collection"])
188 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
189 self.assertIn(noNonRunCollectionsMsg, result.stdout)
191 # Add a couple chained collections
192 for parent, child in (
193 ("chained-run-1", "ingest/run"),
194 ("chained-run-2", "ingest/run"),
195 ):
196 result = self.runner.invoke(
197 butlerCli,
198 ["collection-chain", self.root, parent, child],
199 )
200 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
202 # Test aborting a removal
204 result = self.runner.invoke(
205 butlerCli,
206 ["remove-collections", self.root, "chained-run-1"],
207 input="no",
208 )
209 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
210 self.assertIn(abortedMsg, result.stdout)
212 # Remove with --no-confirm, it's expected to run silently.
214 result = self.runner.invoke(
215 butlerCli, ["remove-collections", self.root, "chained-run-1", "--no-confirm"]
216 )
217 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
218 self.assertIn(removedCollectionsMsg, result.stdout)
219 self.assertIn("chained-run-1", result.stdout)
221 # verify chained-run-1 was removed:
223 butler = Butler.from_config(self.root)
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)
279if __name__ == "__main__":
280 unittest.main()