Coverage for tests/test_cliCmdQueryCollections.py: 28%

97 statements  

« prev     ^ index     » next       coverage.py v7.2.7, created at 2023-07-12 10:56 -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 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 query-collections command. 

23""" 

24 

25import os 

26import unittest 

27 

28from astropy.table import Table 

29from lsst.daf.butler import Butler, CollectionType 

30from lsst.daf.butler.cli.butler import cli 

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

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

33from lsst.daf.butler.script import queryCollections 

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

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

36from numpy import array 

37 

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

39 

40 

41class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

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

43 

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

45 

46 @staticmethod 

47 def defaultExpected(): 

48 return dict( 

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

50 ) 

51 

52 @staticmethod 

53 def command(): 

54 return query_collections 

55 

56 def test_minimal(self): 

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

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

59 

60 def test_all(self): 

61 """Test all parameters""" 

62 self.run_test( 

63 [ 

64 "query-collections", 

65 "here", 

66 "foo*", 

67 "--collection-type", 

68 "TAGGED", 

69 "--collection-type", 

70 "RUN", 

71 "--chains", 

72 "TABLE", 

73 ], 

74 self.makeExpected( 

75 repo="here", 

76 glob=("foo*",), 

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

78 chains="TABLE", 

79 ), 

80 ) 

81 

82 

83class QueryCollectionsScriptTest(ButlerTestHelper, unittest.TestCase): 

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

85 

86 def setUp(self): 

87 self.runner = LogCliRunner() 

88 

89 def testGetCollections(self): 

90 run = "ingest/run" 

91 tag = "tag" 

92 with self.runner.isolated_filesystem(): 

93 butlerCfg = Butler.makeRepo("here") 

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

95 butler = Butler(butlerCfg, run=run, collections=[tag], writeable=True) 

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

97 

98 # Verify collections that were created are found by 

99 # query-collections. 

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

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

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

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

104 

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

106 # name matches with the specified pattern are returned. 

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

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

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

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

111 

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

113 # that type are returned. 

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

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

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

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

118 

119 

120class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

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

122 

123 def setUp(self): 

124 self.runner = LogCliRunner() 

125 

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

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

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

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

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

131 

132 def testChained(self): 

133 with self.runner.isolated_filesystem(): 

134 # Create a butler and add some chained collections: 

135 butlerCfg = Butler.makeRepo("here") 

136 

137 butler1 = Butler(butlerCfg, writeable=True) 

138 

139 # Replace datastore functions with mocks: 

140 DatastoreMock.apply(butler1) 

141 

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

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

144 registry1 = butler1.registry 

145 registry1.registerRun("run1") 

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

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

148 

149 # Create the collection chain 

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

151 self.assertChain( 

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

153 ) 

154 

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

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

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

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

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

160 # the command line interface. 

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

162 

163 expected = Table( 

164 array( 

165 ( 

166 ("calibration1", "CALIBRATION"), 

167 ("chain1", "CHAINED"), 

168 (" tag1", "TAGGED"), 

169 (" run1", "RUN"), 

170 (" chain2", "CHAINED"), 

171 (" calibration1", "CALIBRATION"), 

172 (" run1", "RUN"), 

173 ("chain2", "CHAINED"), 

174 (" calibration1", "CALIBRATION"), 

175 (" run1", "RUN"), 

176 ("imported_g", "RUN"), 

177 ("imported_r", "RUN"), 

178 ("run1", "RUN"), 

179 ("tag1", "TAGGED"), 

180 ) 

181 ), 

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

183 ) 

184 self.assertAstropyTablesEqual(table, expected) 

185 

186 # Test table with inverse == True 

187 table = queryCollections( 

188 "here", 

189 glob=(), 

190 collection_type=CollectionType.all(), 

191 chains="INVERSE-TREE", 

192 ) 

193 expected = Table( 

194 array( 

195 ( 

196 ("calibration1", "CALIBRATION"), 

197 (" chain2", "CHAINED"), 

198 (" chain1", "CHAINED"), 

199 ("chain1", "CHAINED"), 

200 ("chain2", "CHAINED"), 

201 (" chain1", "CHAINED"), 

202 ("imported_g", "RUN"), 

203 ("imported_r", "RUN"), 

204 ("run1", "RUN"), 

205 (" chain1", "CHAINED"), 

206 (" chain2", "CHAINED"), 

207 (" chain1", "CHAINED"), 

208 ("tag1", "TAGGED"), 

209 (" chain1", "CHAINED"), 

210 ) 

211 ), 

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

213 ) 

214 self.assertAstropyTablesEqual(table, expected) 

215 

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

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

218 expected = Table( 

219 array( 

220 ( 

221 ("calibration1", "CALIBRATION", ""), 

222 ("chain1", "CHAINED", "chain2"), 

223 ("", "", "run1"), 

224 ("", "", "tag1"), 

225 ("chain2", "CHAINED", "calibration1"), 

226 ("", "", "run1"), 

227 ("imported_g", "RUN", ""), 

228 ("imported_r", "RUN", ""), 

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

230 ("tag1", "TAGGED", ""), 

231 ) 

232 ), 

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

234 ) 

235 table = readTable(result.output) 

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

237 

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

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

240 expected = Table( 

241 array( 

242 ( 

243 ("calibration1", "CALIBRATION", "chain2"), 

244 ("chain1", "CHAINED", ""), 

245 ("chain2", "CHAINED", "chain1"), 

246 ("imported_g", "RUN", ""), 

247 ("imported_r", "RUN", ""), 

248 ("run1", "RUN", "chain1"), 

249 ("", "", "chain2"), 

250 ("tag1", "TAGGED", "chain1"), 

251 ) 

252 ), 

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

254 ) 

255 table = readTable(result.output) 

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

257 

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

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

260 expected = Table( 

261 array( 

262 ( 

263 ("calibration1", "CALIBRATION"), 

264 ("imported_g", "RUN"), 

265 ("imported_r", "RUN"), 

266 ("run1", "RUN"), 

267 ("tag1", "TAGGED"), 

268 ) 

269 ), 

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

271 ) 

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

273 

274 # Add a couple more run collections for chain testing 

275 registry1.registerRun("run2") 

276 registry1.registerRun("run3") 

277 registry1.registerRun("run4") 

278 

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

280 

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

282 

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

284 

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

286 

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

288 

289 self.assertChain( 

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

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

292 ) 

293 

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

295 

296 

297if __name__ == "__main__": 

298 unittest.main()