Coverage for tests/test_cliCmdRemoveCollections.py: 20%

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/>. 

27 

28"""Unit tests for daf_butler CLI prune-collections subcommand.""" 

29 

30import os 

31import unittest 

32from collections.abc import Sequence 

33 

34from astropy.table import Table 

35from numpy import array 

36 

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) 

57 

58TESTDIR = os.path.abspath(os.path.dirname(__file__)) 

59 

60QueryCollectionsRow = tuple[str, str] | tuple[str, str, str] 

61RemoveCollectionRow = tuple[str, str] 

62 

63 

64class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper): 

65 """Test executing remove collection.""" 

66 

67 def setUp(self): 

68 self.runner = LogCliRunner() 

69 

70 self.root = makeTestTempDir(TESTDIR) 

71 self.testRepo = MetricTestRepo( 

72 self.root, configFile=os.path.join(TESTDIR, "config/basic/butler.yaml") 

73 ) 

74 

75 def tearDown(self): 

76 removeTestTempDir(self.root) 

77 

78 def _verify_remove( 

79 self, 

80 collection: str, 

81 before_rows: Sequence[QueryCollectionsRow], 

82 remove_rows: Sequence[RemoveCollectionRow], 

83 after_rows: Sequence[QueryCollectionsRow], 

84 ): 

85 """Remove collections, with verification that expected collections are 

86 present before removing, that the command reports expected collections 

87 to be removed, and that expected collections are present after removal. 

88 

89 Parameters 

90 ---------- 

91 collection : `str` 

92 The name of the collection, or glob pattern for collections, to 

93 remove. 

94 before_rows : `Sequence` [ `QueryCollectionsRow` ] 

95 The rows that should be in the table returned by query-collections 

96 before removing the collection. 

97 remove_rows : `Sequence` [ `RemoveCollectionRow` ] 

98 The rows that should be in the "will remove" table while removing 

99 collections. 

100 after_rows : `Sequence` [ `QueryCollectionsRow` ] 

101 The rows that should be in the table returned by query-collections 

102 after removing the collection. 

103 """ 

104 

105 def _query_collection_column_names(rows): 

106 # If there is a chained collection in the table then there is a 

107 # definition column, otherwise there is only the name and type 

108 # columns. 

109 if len(rows[0]) == 2: 

110 return ("Name", "Type") 

111 elif len(rows[0]) == 3: 

112 return ("Name", "Type", "Children") 

113 else: 

114 raise RuntimeError(f"Unhandled column count: {len(rows[0])}") 

115 

116 result = self.runner.invoke(butlerCli, ["query-collections", self.root, "--chains", "TABLE"]) 

117 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

118 expected = Table(array(before_rows), names=_query_collection_column_names(before_rows)) 

119 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True) 

120 

121 removal = removeCollections(repo=self.root, collection=collection, remove_from_parents=False) 

122 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

123 expected = Table(array(remove_rows), names=("Collection", "Collection Type")) 

124 self.assertAstropyTablesEqual(removal.removeCollectionsTable, expected) 

125 removal.onConfirmation() 

126 

127 result = self.runner.invoke(butlerCli, ["query-collections", self.root]) 

128 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

129 expected = Table(array(after_rows), names=_query_collection_column_names(after_rows)) 

130 self.assertAstropyTablesEqual(readTable(result.output), expected, unorderedRows=True) 

131 

132 def testRemoveScript(self): 

133 """Test removing collections. 

134 

135 Combining several tests into one case allows us to reuse the test repo, 

136 which saves execution time. 

137 """ 

138 # Test wildcard with chained collections: 

139 

140 # Add a couple chained collections 

141 for parent, child in ( 

142 ("chained-run-1", "ingest/run"), 

143 ("chained-run-2", "ingest/run"), 

144 ): 

145 result = self.runner.invoke( 

146 butlerCli, 

147 ["collection-chain", self.root, parent, child], 

148 ) 

149 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

150 

151 self._verify_remove( 

152 collection="chained-run-*", 

153 before_rows=( 

154 ("chained-run-1", "CHAINED", "ingest/run"), 

155 ("chained-run-2", "CHAINED", "ingest/run"), 

156 ("ingest", "TAGGED", ""), 

157 ("ingest/run", "RUN", ""), 

158 ), 

159 remove_rows=( 

160 ("chained-run-1", "CHAINED"), 

161 ("chained-run-2", "CHAINED"), 

162 ), 

163 after_rows=( 

164 ("ingest", "TAGGED"), 

165 ("ingest/run", "RUN"), 

166 ), 

167 ) 

168 

169 # Test a single tagged collection: 

170 

171 self._verify_remove( 

172 collection="ingest", 

173 before_rows=( 

174 ("ingest", "TAGGED"), 

175 ("ingest/run", "RUN"), 

176 ), 

177 remove_rows=(("ingest", "TAGGED"),), 

178 after_rows=(("ingest/run", "RUN"),), 

179 ) 

180 

181 def testRemoveCmd(self): 

182 """Test remove command outputs.""" 

183 # Test expected output with a non-existent collection: 

184 

185 result = self.runner.invoke(butlerCli, ["remove-collections", self.root, "fake_collection"]) 

186 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

187 self.assertIn(noNonRunCollectionsMsg, result.stdout) 

188 

189 # Add a couple chained collections 

190 for parent, child in ( 

191 ("chained-run-1", "ingest/run"), 

192 ("chained-run-2", "ingest/run"), 

193 ): 

194 result = self.runner.invoke( 

195 butlerCli, 

196 ["collection-chain", self.root, parent, child], 

197 ) 

198 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

199 

200 # Test aborting a removal 

201 

202 result = self.runner.invoke( 

203 butlerCli, 

204 ["remove-collections", self.root, "chained-run-1"], 

205 input="no", 

206 ) 

207 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

208 self.assertIn(abortedMsg, result.stdout) 

209 

210 # Remove with --no-confirm, it's expected to run silently. 

211 

212 result = self.runner.invoke( 

213 butlerCli, ["remove-collections", self.root, "chained-run-1", "--no-confirm"] 

214 ) 

215 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

216 self.assertIn(removedCollectionsMsg, result.stdout) 

217 self.assertIn("chained-run-1", result.stdout) 

218 

219 # verify chained-run-1 was removed: 

220 

221 butler = Butler.from_config(self.root) 

222 collections = butler.registry.queryCollections( 

223 collectionTypes=frozenset( 

224 ( 

225 CollectionType.RUN, 

226 CollectionType.TAGGED, 

227 CollectionType.CHAINED, 

228 CollectionType.CALIBRATION, 

229 ) 

230 ), 

231 ) 

232 self.assertCountEqual(["ingest/run", "ingest", "chained-run-2"], collections) 

233 

234 # verify chained-run-2 can be removed with prompting and expected CLI 

235 # output 

236 

237 result = self.runner.invoke( 

238 butlerCli, 

239 ["remove-collections", self.root, "chained-run-2"], 

240 input="yes", 

241 ) 

242 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

243 self.assertIn(willRemoveCollectionMsg, result.stdout) 

244 self.assertIn("chained-run-2 CHAINED", result.stdout) 

245 

246 # try to remove a run table, check for the "can not remove run" message 

247 

248 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child]) 

249 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

250 result = self.runner.invoke( 

251 # removes run-chain (chained collection), but can not remove the 

252 # run collection, and emits a message that says so. 

253 butlerCli, 

254 ["remove-collections", self.root, "*run*"], 

255 input="yes", 

256 ) 

257 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

258 self.assertIn(canNotRemoveFoundRuns, result.stdout) 

259 self.assertIn("ingest/run", result.stdout) 

260 

261 # try to remove a run table with --no-confirm, check for the "did not 

262 # remove run" message 

263 

264 result = self.runner.invoke(butlerCli, ["collection-chain", self.root, "run-chain", child]) 

265 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

266 result = self.runner.invoke( 

267 # removes run-chain (chained collection), but can not remove the 

268 # run collection, and emits a message that says so. 

269 butlerCli, 

270 ["remove-collections", self.root, "*run*", "--no-confirm"], 

271 ) 

272 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

273 self.assertIn(didNotRemoveFoundRuns, result.stdout) 

274 self.assertIn("ingest/run", result.stdout) 

275 

276 def testRemoveFromParents(self) -> None: 

277 butler = Butler(self.root, writeable=True) 

278 butler.collections.register("tag1", CollectionType.TAGGED) 

279 butler.collections.register("tag2", CollectionType.TAGGED) 

280 butler.collections.register("chain1", CollectionType.CHAINED) 

281 butler.collections.register("chain2", CollectionType.CHAINED) 

282 butler.collections.register("chain3", CollectionType.CHAINED) 

283 butler.collections.redefine_chain("chain1", ["tag1", "tag2", "chain3"]) 

284 butler.collections.redefine_chain("chain2", ["tag1"]) 

285 

286 # Make sure the printed output is correct. 

287 removal = removeCollections(repo=self.root, collection="tag*", remove_from_parents=True) 

288 table = [tuple(row) for row in removal.removeChainsTable] 

289 self.assertEqual(table, [("tag1", "chain1"), ("", "chain2"), ("tag2", "chain1")]) 

290 result = self.runner.invoke( 

291 butlerCli, 

292 ["remove-collections", "--remove-from-parents", self.root, "tag*"], 

293 input="yes", 

294 ) 

295 self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

296 self.assertIn(willRemoveCollectionChainsMsg, result.stdout) 

297 # Make sure the collections are deleted as expected. 

298 self.assertEqual( 

299 sorted(butler.collections.query("*")), ["chain1", "chain2", "chain3", "ingest", "ingest/run"] 

300 ) 

301 self.assertEqual(butler.collections.get_info("chain1").children, ("chain3",)) 

302 

303 

304if __name__ == "__main__": 

305 unittest.main()