Coverage for tests/test_cliCmdRemoveCollections.py: 24%
81 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-21 09:55 +0000
« prev ^ index » next coverage.py v7.2.7, created at 2023-07-21 09:55 +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-collections subcommand.
23"""
25import os
26import unittest
27from collections.abc import Sequence
29from astropy.table import Table
30from lsst.daf.butler import Butler, CollectionType
31from lsst.daf.butler.cli.butler import cli as butlerCli
32from lsst.daf.butler.cli.cmd._remove_collections import (
33 abortedMsg,
34 canNotRemoveFoundRuns,
35 didNotRemoveFoundRuns,
36 noNonRunCollectionsMsg,
37 removedCollectionsMsg,
38 willRemoveCollectionMsg,
39)
40from lsst.daf.butler.cli.utils import LogCliRunner, clickResultMsg
41from lsst.daf.butler.script.removeCollections import removeCollections
42from lsst.daf.butler.tests.utils import (
43 ButlerTestHelper,
44 MetricTestRepo,
45 makeTestTempDir,
46 readTable,
47 removeTestTempDir,
48)
49from numpy import array
51TESTDIR = os.path.abspath(os.path.dirname(__file__))
53QueryCollectionsRow = tuple[str, str] | tuple[str, str, str]
54RemoveCollectionRow = tuple[str, str]
57class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper):
58 """Test executing remove collection."""
60 def setUp(self):
61 self.runner = LogCliRunner()
63 self.root = makeTestTempDir(TESTDIR)
64 self.testRepo = MetricTestRepo(
65 self.root, configFile=os.path.join(TESTDIR, "config/basic/butler.yaml")
66 )
68 def tearDown(self):
69 removeTestTempDir(self.root)
71 def _verify_remove(
72 self,
73 collection: str,
74 before_rows: Sequence[QueryCollectionsRow],
75 remove_rows: Sequence[RemoveCollectionRow],
76 after_rows: Sequence[QueryCollectionsRow],
77 ):
78 """Remove collections, with verification that expected collections are
79 present before removing, that the command reports expected collections
80 to be removed, and that expected collections are present after removal.
82 Parameters
83 ----------
84 collection : `str`
85 The name of the collection, or glob pattern for collections, to
86 remove.
87 before_rows : `Sequence` [ `QueryCollectionsRow` ]
88 The rows that should be in the table returned by query-collections
89 before removing the collection.
90 remove_rows : `Sequence` [ `RemoveCollectionRow` ]
91 The rows that should be in the "will remove" table while removing
92 collections.
93 after_rows : `Sequence` [ `QueryCollectionsRow` ]
94 The rows that should be in the table returned by query-collections
95 after removing the collection.
96 """
98 def _query_collection_column_names(rows):
99 # If there is a chained collection in the table then there is a
100 # definition column, otherwise there is only the name and type
101 # columns.
102 if len(rows[0]) == 2:
103 return ("Name", "Type")
104 elif len(rows[0]) == 3:
105 return ("Name", "Type", "Children")
106 else:
107 raise RuntimeError(f"Unhandled column count: {len(rows[0])}")
109 result = self.runner.invoke(butlerCli, ["query-collections", self.root, "--chains", "TABLE"])
110 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
111 expected = Table(array(before_rows), names=_query_collection_column_names(before_rows))
112 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
114 removal = removeCollections(
115 repo=self.root,
116 collection=collection,
117 )
118 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
119 expected = Table(array(remove_rows), names=("Collection", "Collection Type"))
120 self.assertAstropyTablesEqual(removal.removeCollectionsTable, expected)
121 removal.onConfirmation()
123 result = self.runner.invoke(butlerCli, ["query-collections", self.root])
124 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
125 expected = Table(array(after_rows), names=_query_collection_column_names(after_rows))
126 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True)
128 def testRemoveScript(self):
129 """Test removing collections.
131 Combining several tests into one case allows us to reuse the test repo,
132 which saves execution time.
133 """
134 # Test wildcard with chained collections:
136 # Add a couple chained collections
137 for parent, child in (
138 ("chained-run-1", "ingest/run"),
139 ("chained-run-2", "ingest/run"),
140 ):
141 result = self.runner.invoke(
142 butlerCli,
143 ["collection-chain", self.root, parent, child],
144 )
145 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
147 self._verify_remove(
148 collection="chained-run-*",
149 before_rows=(
150 ("chained-run-1", "CHAINED", "ingest/run"),
151 ("chained-run-2", "CHAINED", "ingest/run"),
152 ("ingest", "TAGGED", ""),
153 ("ingest/run", "RUN", ""),
154 ),
155 remove_rows=(
156 ("chained-run-1", "CHAINED"),
157 ("chained-run-2", "CHAINED"),
158 ),
159 after_rows=(
160 ("ingest", "TAGGED"),
161 ("ingest/run", "RUN"),
162 ),
163 )
165 # Test a single tagged collection:
167 self._verify_remove(
168 collection="ingest",
169 before_rows=(
170 ("ingest", "TAGGED"),
171 ("ingest/run", "RUN"),
172 ),
173 remove_rows=(("ingest", "TAGGED"),),
174 after_rows=(("ingest/run", "RUN"),),
175 )
177 def testRemoveCmd(self):
178 """Test remove command outputs."""
179 # Test expected output with a non-existent collection:
181 result = self.runner.invoke(butlerCli, ["remove-collections", self.root, "fake_collection"])
182 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
183 self.assertIn(noNonRunCollectionsMsg, result.stdout)
185 # Add a couple chained collections
186 for parent, child in (
187 ("chained-run-1", "ingest/run"),
188 ("chained-run-2", "ingest/run"),
189 ):
190 result = self.runner.invoke(
191 butlerCli,
192 ["collection-chain", self.root, parent, child],
193 )
194 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
196 # Test aborting a removal
198 result = self.runner.invoke(
199 butlerCli,
200 ["remove-collections", self.root, "chained-run-1"],
201 input="no",
202 )
203 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
204 self.assertIn(abortedMsg, result.stdout)
206 # Remove with --no-confirm, it's expected to run silently.
208 result = self.runner.invoke(
209 butlerCli, ["remove-collections", self.root, "chained-run-1", "--no-confirm"]
210 )
211 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
212 self.assertIn(removedCollectionsMsg, result.stdout)
213 self.assertIn("chained-run-1", result.stdout)
215 # verify chained-run-1 was removed:
217 butler = Butler(self.root)
218 collections = butler.registry.queryCollections(
219 collectionTypes=frozenset(
220 (
221 CollectionType.RUN,
222 CollectionType.TAGGED,
223 CollectionType.CHAINED,
224 CollectionType.CALIBRATION,
225 )
226 ),
227 )
228 self.assertCountEqual(["ingest/run", "ingest", "chained-run-2"], collections)
230 # verify chained-run-2 can be removed with prompting and expected CLI
231 # output
233 result = self.runner.invoke(
234 butlerCli,
235 ["remove-collections", self.root, "chained-run-2"],
236 input="yes",
237 )
238 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
239 self.assertIn(willRemoveCollectionMsg, result.stdout)
240 self.assertIn("chained-run-2 CHAINED", result.stdout)
242 # try to remove a run table, check for the "can not remove run" message
244 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child])
245 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
246 result = self.runner.invoke(
247 # removes run-chain (chained collection), but can not remove the
248 # run collection, and emits a message that says so.
249 butlerCli,
250 ["remove-collections", self.root, "*run*"],
251 input="yes",
252 )
253 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
254 self.assertIn(canNotRemoveFoundRuns, result.stdout)
255 self.assertIn("ingest/run", result.stdout)
257 # try to remove a run table with --no-confirm, check for the "did not
258 # remove run" message
260 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child])
261 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
262 result = self.runner.invoke(
263 # removes run-chain (chained collection), but can not remove the
264 # run collection, and emits a message that says so.
265 butlerCli,
266 ["remove-collections", self.root, "*run*", "--no-confirm"],
267 )
268 self.assertEqual(result.exit_code, 0, clickResultMsg(result))
269 self.assertIn(didNotRemoveFoundRuns, result.stdout)
270 self.assertIn("ingest/run", result.stdout)
273if __name__ == "__main__":
274 unittest.main()