Coverage for tests/test_cliCmdRemoveCollections.py: 24%

81 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-16 02:58 -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/>. 

27 

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

29""" 

30 

31import os 

32import unittest 

33from collections.abc import Sequence 

34 

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 

56 

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

58 

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

60RemoveCollectionRow = tuple[str, str] 

61 

62 

63class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper): 

64 """Test executing remove collection.""" 

65 

66 def setUp(self): 

67 self.runner = LogCliRunner() 

68 

69 self.root = makeTestTempDir(TESTDIR) 

70 self.testRepo = MetricTestRepo( 

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

72 ) 

73 

74 def tearDown(self): 

75 removeTestTempDir(self.root) 

76 

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. 

87 

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 """ 

103 

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])}") 

114 

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) 

119 

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() 

128 

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) 

133 

134 def testRemoveScript(self): 

135 """Test removing collections. 

136 

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: 

141 

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)) 

152 

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 ) 

170 

171 # Test a single tagged collection: 

172 

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 ) 

182 

183 def testRemoveCmd(self): 

184 """Test remove command outputs.""" 

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

186 

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) 

190 

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)) 

201 

202 # Test aborting a removal 

203 

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) 

211 

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

213 

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) 

220 

221 # verify chained-run-1 was removed: 

222 

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) 

235 

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

237 # output 

238 

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) 

247 

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

249 

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) 

262 

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

264 # remove run" message 

265 

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) 

277 

278 

279if __name__ == "__main__": 

280 unittest.main()