Coverage for tests/test_cliCmdQueryCollections.py: 30%

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 query-collections command. 

29""" 

30 

31import os 

32import unittest 

33 

34from astropy.table import Table 

35from lsst.daf.butler import Butler, CollectionType 

36from lsst.daf.butler.cli.butler import cli 

37from lsst.daf.butler.cli.cmd import query_collections 

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

39from lsst.daf.butler.script import queryCollections 

40from lsst.daf.butler.tests import CliCmdTestBase, DatastoreMock 

41from lsst.daf.butler.tests.utils import ButlerTestHelper, readTable 

42from numpy import array 

43 

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

45 

46 

47class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

48 """Test the query-collections command-line.""" 

49 

50 mockFuncName = "lsst.daf.butler.cli.cmd.commands.script.queryCollections" 

51 

52 @staticmethod 

53 def defaultExpected(): 

54 return dict( 

55 repo=None, collection_type=tuple(CollectionType.__members__.values()), chains="TABLE", glob=() 

56 ) 

57 

58 @staticmethod 

59 def command(): 

60 return query_collections 

61 

62 def test_minimal(self): 

63 """Test only required parameters, and omit optional parameters.""" 

64 self.run_test(["query-collections", "here", "--chains", "TABLE"], self.makeExpected(repo="here")) 

65 

66 def test_all(self): 

67 """Test all parameters""" 

68 self.run_test( 

69 [ 

70 "query-collections", 

71 "here", 

72 "foo*", 

73 "--collection-type", 

74 "TAGGED", 

75 "--collection-type", 

76 "RUN", 

77 "--chains", 

78 "TABLE", 

79 ], 

80 self.makeExpected( 

81 repo="here", 

82 glob=("foo*",), 

83 collection_type=(CollectionType.TAGGED, CollectionType.RUN), 

84 chains="TABLE", 

85 ), 

86 ) 

87 

88 

89class QueryCollectionsScriptTest(ButlerTestHelper, unittest.TestCase): 

90 """Test the query-collections script interface.""" 

91 

92 def setUp(self): 

93 self.runner = LogCliRunner() 

94 

95 def testGetCollections(self): 

96 run = "ingest/run" 

97 tag = "tag" 

98 with self.runner.isolated_filesystem(): 

99 butlerCfg = Butler.makeRepo("here") 

100 # the purpose of this call is to create some collections 

101 butler = Butler.from_config(butlerCfg, run=run, collections=[tag], writeable=True) 

102 butler.registry.registerCollection(tag, CollectionType.TAGGED) 

103 

104 # Verify collections that were created are found by 

105 # query-collections. 

106 result = self.runner.invoke(cli, ["query-collections", "here"]) 

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

108 expected = Table((("ingest/run", "tag"), ("RUN", "TAGGED")), names=("Name", "Type")) 

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

110 

111 # Verify that with a glob argument, that only collections whose 

112 # name matches with the specified pattern are returned. 

113 result = self.runner.invoke(cli, ["query-collections", "here", "t*"]) 

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

115 expected = Table((("tag",), ("TAGGED",)), names=("Name", "Type")) 

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

117 

118 # Verify that with a collection type argument, only collections of 

119 # that type are returned. 

120 result = self.runner.invoke(cli, ["query-collections", "here", "--collection-type", "RUN"]) 

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

122 expected = Table((("ingest/run",), ("RUN",)), names=("Name", "Type")) 

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

124 

125 

126class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

127 """Test the collection-chain command-line interface.""" 

128 

129 def setUp(self): 

130 self.runner = LogCliRunner() 

131 

132 def assertChain(self, args: list[str], expected: str): 

133 """Run collection-chain and check the expected result""" 

134 result = self.runner.invoke(cli, ["collection-chain", "here", *args]) 

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

136 self.assertEqual(result.output.strip(), expected, clickResultMsg(result)) 

137 

138 def testChained(self): 

139 with self.runner.isolated_filesystem(): 

140 # Create a butler and add some chained collections: 

141 butlerCfg = Butler.makeRepo("here") 

142 

143 butler1 = Butler.from_config(butlerCfg, writeable=True) 

144 

145 # Replace datastore functions with mocks: 

146 DatastoreMock.apply(butler1) 

147 

148 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

149 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml")) 

150 registry1 = butler1.registry 

151 registry1.registerRun("run1") 

152 registry1.registerCollection("tag1", CollectionType.TAGGED) 

153 registry1.registerCollection("calibration1", CollectionType.CALIBRATION) 

154 

155 # Create the collection chain 

156 self.assertChain(["chain2", "calibration1", "run1"], "[calibration1, run1]") 

157 self.assertChain( 

158 ["--mode", "redefine", "chain1", "tag1", "run1", "chain2"], "[tag1, run1, chain2]" 

159 ) 

160 

161 # Use the script function to test the query-collections TREE 

162 # option, because the astropy.table.Table.read method, which we are 

163 # using for verification elsewhere in this file, seems to strip 

164 # leading whitespace from columns. This makes it impossible to test 

165 # the nested TREE output of the query-collections subcommand from 

166 # the command line interface. 

167 table = queryCollections("here", glob=(), collection_type=CollectionType.all(), chains="TREE") 

168 

169 expected = Table( 

170 array( 

171 ( 

172 ("calibration1", "CALIBRATION"), 

173 ("chain1", "CHAINED"), 

174 (" tag1", "TAGGED"), 

175 (" run1", "RUN"), 

176 (" chain2", "CHAINED"), 

177 (" calibration1", "CALIBRATION"), 

178 (" run1", "RUN"), 

179 ("chain2", "CHAINED"), 

180 (" calibration1", "CALIBRATION"), 

181 (" run1", "RUN"), 

182 ("imported_g", "RUN"), 

183 ("imported_r", "RUN"), 

184 ("run1", "RUN"), 

185 ("tag1", "TAGGED"), 

186 ) 

187 ), 

188 names=("Name", "Type"), 

189 ) 

190 self.assertAstropyTablesEqual(table, expected) 

191 

192 # Test table with inverse == True 

193 table = queryCollections( 

194 "here", 

195 glob=(), 

196 collection_type=CollectionType.all(), 

197 chains="INVERSE-TREE", 

198 ) 

199 expected = Table( 

200 array( 

201 ( 

202 ("calibration1", "CALIBRATION"), 

203 (" chain2", "CHAINED"), 

204 (" chain1", "CHAINED"), 

205 ("chain1", "CHAINED"), 

206 ("chain2", "CHAINED"), 

207 (" chain1", "CHAINED"), 

208 ("imported_g", "RUN"), 

209 ("imported_r", "RUN"), 

210 ("run1", "RUN"), 

211 (" chain1", "CHAINED"), 

212 (" chain2", "CHAINED"), 

213 (" chain1", "CHAINED"), 

214 ("tag1", "TAGGED"), 

215 (" chain1", "CHAINED"), 

216 ) 

217 ), 

218 names=("Name", "Type"), 

219 ) 

220 self.assertAstropyTablesEqual(table, expected) 

221 

222 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "TABLE"]) 

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

224 expected = Table( 

225 array( 

226 ( 

227 ("calibration1", "CALIBRATION", ""), 

228 ("chain1", "CHAINED", "tag1"), 

229 ("", "", "run1"), 

230 ("", "", "chain2"), 

231 ("chain2", "CHAINED", "calibration1"), 

232 ("", "", "run1"), 

233 ("imported_g", "RUN", ""), 

234 ("imported_r", "RUN", ""), 

235 ("run1", "RUN", ""), 

236 ("tag1", "TAGGED", ""), 

237 ) 

238 ), 

239 names=("Name", "Type", "Children"), 

240 ) 

241 table = readTable(result.output) 

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

243 

244 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "INVERSE-TABLE"]) 

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

246 expected = Table( 

247 array( 

248 ( 

249 ("calibration1", "CALIBRATION", "chain2"), 

250 ("chain1", "CHAINED", ""), 

251 ("chain2", "CHAINED", "chain1"), 

252 ("imported_g", "RUN", ""), 

253 ("imported_r", "RUN", ""), 

254 ("run1", "RUN", "chain1"), 

255 ("", "", "chain2"), 

256 ("tag1", "TAGGED", "chain1"), 

257 ) 

258 ), 

259 names=("Name", "Type", "Parents"), 

260 ) 

261 table = readTable(result.output) 

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

263 

264 result = self.runner.invoke(cli, ["query-collections", "here", "--chains", "FLATTEN"]) 

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

266 expected = Table( 

267 array( 

268 ( 

269 ("calibration1", "CALIBRATION"), 

270 ("imported_g", "RUN"), 

271 ("imported_r", "RUN"), 

272 ("run1", "RUN"), 

273 ("tag1", "TAGGED"), 

274 ) 

275 ), 

276 names=("Name", "Type"), 

277 ) 

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

279 

280 # Add a couple more run collections for chain testing 

281 registry1.registerRun("run2") 

282 registry1.registerRun("run3") 

283 registry1.registerRun("run4") 

284 

285 self.assertChain(["--mode", "pop", "chain1"], "[run1, chain2]") 

286 

287 self.assertChain(["--mode", "extend", "chain1", "run2", "run3"], "[run1, chain2, run2, run3]") 

288 

289 self.assertChain(["--mode", "remove", "chain1", "chain2", "run2"], "[run1, run3]") 

290 

291 self.assertChain(["--mode", "prepend", "chain1", "chain2", "run2"], "[chain2, run2, run1, run3]") 

292 

293 self.assertChain(["--mode", "pop", "chain1", "1", "3"], "[chain2, run1]") 

294 

295 self.assertChain( 

296 ["--mode", "redefine", "chain1", "chain2", "run2", "run3,run4", "--flatten"], 

297 "[calibration1, run1, run2, run3, run4]", 

298 ) 

299 

300 self.assertChain(["--mode", "pop", "chain1", "--", "-1", "-3"], "[calibration1, run1, run3]") 

301 

302 # Out-of-bounds index 

303 result = self.runner.invoke(cli, ["collection-chain", "here", "--mode", "pop", "chain1", "10"]) 

304 self.assertEqual(result.exit_code, 1) 

305 

306 

307if __name__ == "__main__": 

308 unittest.main()