Coverage for tests/test_cliCmdRemoveCollections.py: 26%

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

83 statements  

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

21 

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

23""" 

24 

25import os 

26import unittest 

27from typing import Sequence, Tuple, Union 

28 

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 

50 

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

52 

53QueryCollectionsRow = Union[Tuple[str, str], Tuple[str, str, str]] 

54RemoveCollectionRow = Tuple[str, str] 

55 

56 

57class RemoveCollectionTest(unittest.TestCase, ButlerTestHelper): 

58 """Test executing remove collection.""" 

59 

60 def setUp(self): 

61 self.runner = LogCliRunner() 

62 

63 self.root = makeTestTempDir(TESTDIR) 

64 self.testRepo = MetricTestRepo( 

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

66 ) 

67 

68 def tearDown(self): 

69 removeTestTempDir(self.root) 

70 

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. 

81 

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

97 

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

108 

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) 

113 

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

122 

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) 

127 

128 def testRemoveScript(self): 

129 """Test removing collections. 

130 

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

132 which saves execution time. 

133 """ 

134 

135 # Test wildcard with chained collections: 

136 

137 # Add a couple chained collections 

138 for parent, child in ( 

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

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

141 ): 

142 result = self.runner.invoke( 

143 butlerCli, 

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

145 ) 

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

147 

148 self._verify_remove( 

149 collection="chained-run-*", 

150 before_rows=( 

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

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

153 ("ingest", "TAGGED", ""), 

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

155 ), 

156 remove_rows=( 

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

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

159 ), 

160 after_rows=( 

161 ("ingest", "TAGGED"), 

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

163 ), 

164 ) 

165 

166 # Test a single tagged collection: 

167 

168 self._verify_remove( 

169 collection="ingest", 

170 before_rows=( 

171 ("ingest", "TAGGED"), 

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

173 ), 

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

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

176 ) 

177 

178 def testRemoveCmd(self): 

179 """Test remove command outputs.""" 

180 

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

182 

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

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

185 self.assertIn(noNonRunCollectionsMsg, result.stdout) 

186 

187 # Add a couple chained collections 

188 for parent, child in ( 

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

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

191 ): 

192 result = self.runner.invoke( 

193 butlerCli, 

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

195 ) 

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

197 

198 # Test aborting a removal 

199 

200 result = self.runner.invoke( 

201 butlerCli, 

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

203 input="no", 

204 ) 

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

206 self.assertIn(abortedMsg, result.stdout) 

207 

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

209 

210 result = self.runner.invoke( 

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

212 ) 

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

214 self.assertIn(removedCollectionsMsg, result.stdout) 

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

216 

217 # verify chained-run-1 was removed: 

218 

219 butler = Butler(self.root) 

220 collections = butler.registry.queryCollections( 

221 collectionTypes=frozenset( 

222 ( 

223 CollectionType.RUN, 

224 CollectionType.TAGGED, 

225 CollectionType.CHAINED, 

226 CollectionType.CALIBRATION, 

227 ) 

228 ), 

229 ) 

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

231 

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

233 # output 

234 

235 result = self.runner.invoke( 

236 butlerCli, 

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

238 input="yes", 

239 ) 

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

241 self.assertIn(willRemoveCollectionMsg, result.stdout) 

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

243 

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

245 

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

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

248 result = self.runner.invoke( 

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

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

251 butlerCli, 

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

253 input="yes", 

254 ) 

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

256 self.assertIn(canNotRemoveFoundRuns, result.stdout) 

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

258 

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

260 # remove run" message 

261 

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

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

264 result = self.runner.invoke( 

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

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

267 butlerCli, 

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

269 ) 

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

271 self.assertIn(didNotRemoveFoundRuns, result.stdout) 

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

273 

274 

275if __name__ == "__main__": 275 ↛ 276line 275 didn't jump to line 276, because the condition on line 275 was never true

276 unittest.main()