Coverage for tests/test_cliCmdRemoveCollections.py: 24%

81 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-08-05 01:26 +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/>. 

21 

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

23""" 

24 

25import os 

26import unittest 

27from collections.abc import Sequence 

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 = 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 # Test wildcard with chained collections: 

135 

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

146 

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 ) 

164 

165 # Test a single tagged collection: 

166 

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 ) 

176 

177 def testRemoveCmd(self): 

178 """Test remove command outputs.""" 

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

180 

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) 

184 

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

195 

196 # Test aborting a removal 

197 

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) 

205 

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

207 

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) 

214 

215 # verify chained-run-1 was removed: 

216 

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) 

229 

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

231 # output 

232 

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) 

241 

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

243 

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) 

256 

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

258 # remove run" message 

259 

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) 

271 

272 

273if __name__ == "__main__": 

274 unittest.main()