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 

25from astropy.table import Table 

26from typing import Sequence, Tuple, Union 

27from numpy import array 

28import os 

29import unittest 

30 

31from lsst.daf.butler import Butler, CollectionType 

32from lsst.daf.butler.cli.butler import cli as butlerCli 

33from lsst.daf.butler.script.removeCollections import removeCollections 

34from lsst.daf.butler.cli.cmd._remove_collections import ( 

35 abortedMsg, 

36 canNotRemoveFoundRuns, 

37 didNotRemoveFoundRuns, 

38 noNonRunCollectionsMsg, 

39 removedCollectionsMsg, 

40 willRemoveCollectionMsg, 

41) 

42from lsst.daf.butler.cli.utils import clickResultMsg, LogCliRunner 

43from lsst.daf.butler.tests.utils import ( 

44 ButlerTestHelper, 

45 makeTestTempDir, 

46 MetricTestRepo, 

47 readTable, 

48 removeTestTempDir, 

49) 

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 

61 def setUp(self): 

62 self.runner = LogCliRunner() 

63 

64 self.root = makeTestTempDir(TESTDIR) 

65 self.testRepo = MetricTestRepo( 

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

67 ) 

68 

69 def tearDown(self): 

70 removeTestTempDir(self.root) 

71 

72 def _verify_remove( 

73 self, 

74 collection: str, 

75 before_rows: Sequence[QueryCollectionsRow], 

76 remove_rows: Sequence[RemoveCollectionRow], 

77 after_rows: Sequence[QueryCollectionsRow], 

78 ): 

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

80 present before removing, that the command reports expected collections 

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

82 

83 Parameters 

84 ---------- 

85 collection : `str` 

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

87 remove. 

88 before_rows : `Sequence` [ `QueryCollectionsRow` ] 

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

90 before removing the collection. 

91 remove_rows : `Sequence` [ `RemoveCollectionRow` ] 

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

93 collections. 

94 after_rows : `Sequence` [ `QueryCollectionsRow` ] 

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

96 after removing the collection. 

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", "Definition") 

106 else: 

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

108 

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

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) 

113 

114 removal = removeCollections(repo=self.root, collection=collection,) 

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

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

117 self.assertAstropyTablesEqual(removal.removeCollectionsTable, expected) 

118 removal.onConfirmation() 

119 

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

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

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

123 self.assertAstropyTablesEqual(readTable(result.output), expected) 

124 

125 def testRemoveScript(self): 

126 """Test removing collections. 

127 

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

129 which saves execution time. 

130 """ 

131 

132 # Test wildcard with chained collections: 

133 

134 # Add a couple chained collections 

135 for parent, child in ( 

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

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

138 ): 

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

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

141 

142 self._verify_remove( 

143 collection="chained-run-*", 

144 before_rows=( 

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

146 ("ingest", "TAGGED", ""), 

147 ("chained-run-1", "CHAINED", "[ingest/run]"), 

148 ("chained-run-2", "CHAINED", "[ingest/run]"), 

149 ), 

150 remove_rows=(("chained-run-1", "CHAINED"), ("chained-run-2", "CHAINED"),), 

151 after_rows=(("ingest/run", "RUN"), ("ingest", "TAGGED"),), 

152 ) 

153 

154 # Test a single tagged collection: 

155 

156 self._verify_remove( 

157 collection="ingest", 

158 before_rows=(("ingest/run", "RUN"), ("ingest", "TAGGED"),), 

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

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

161 ) 

162 

163 def testRemoveCmd(self): 

164 """Test remove command outputs. 

165 """ 

166 

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

168 

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

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

171 self.assertIn(noNonRunCollectionsMsg, result.stdout) 

172 

173 # Add a couple chained collections 

174 for parent, child in ( 

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

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

177 ): 

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

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

180 

181 # Test aborting a removal 

182 

183 result = self.runner.invoke( 

184 butlerCli, ["remove-collections", self.root, "chained-run-1"], input="no", 

185 ) 

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

187 self.assertIn(abortedMsg, result.stdout) 

188 

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

190 

191 result = self.runner.invoke( 

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

193 ) 

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

195 self.assertIn(removedCollectionsMsg, result.stdout) 

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

197 

198 # verify chained-run-1 was removed: 

199 

200 butler = Butler(self.root) 

201 collections = butler.registry.queryCollections( 

202 collectionTypes=frozenset( 

203 ( 

204 CollectionType.RUN, 

205 CollectionType.TAGGED, 

206 CollectionType.CHAINED, 

207 CollectionType.CALIBRATION, 

208 ) 

209 ), 

210 ) 

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

212 

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

214 # output 

215 

216 result = self.runner.invoke( 

217 butlerCli, ["remove-collections", self.root, "chained-run-2"], input="yes", 

218 ) 

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

220 self.assertIn(willRemoveCollectionMsg, result.stdout) 

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

222 

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

224 

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

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

227 result = self.runner.invoke( 

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

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

230 butlerCli, ["remove-collections", self.root, "*run*"], input="yes", 

231 ) 

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

233 self.assertIn(canNotRemoveFoundRuns, result.stdout) 

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

235 

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

237 # remove run" message 

238 

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

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

241 result = self.runner.invoke( 

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

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

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

245 ) 

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

247 self.assertIn(didNotRemoveFoundRuns, result.stdout) 

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

249 

250 

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

252 unittest.main()