Coverage for tests/test_cliCmdQueryCollections.py: 29%

100 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-01-04 02:04 -0800

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 

27from typing import List 

28 

29from astropy.table import Table 

30from lsst.daf.butler import Butler, CollectionType 

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

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

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

34from lsst.daf.butler.script import queryCollections 

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

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

37from numpy import array 

38 

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

40 

41 

42class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

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 def setUp(self): 

85 self.runner = LogCliRunner() 

86 

87 def testGetCollections(self): 

88 run = "ingest/run" 

89 tag = "tag" 

90 with self.runner.isolated_filesystem(): 

91 butlerCfg = Butler.makeRepo("here") 

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

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

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

95 

96 # Verify collections that were created are found by 

97 # query-collections. 

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

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

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

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

102 

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

104 # name matches with the specified pattern are returned. 

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

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

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

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

109 

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

111 # that type are returned. 

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

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

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

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

116 

117 

118class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

119 def setUp(self): 

120 self.runner = LogCliRunner() 

121 

122 def assertChain(self, args: List[str], expected: str): 

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

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

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

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

127 

128 def testChained(self): 

129 with self.runner.isolated_filesystem(): 

130 

131 # Create a butler and add some chained collections: 

132 butlerCfg = Butler.makeRepo("here") 

133 

134 butler1 = Butler(butlerCfg, writeable=True) 

135 

136 # Replace datastore functions with mocks: 

137 DatastoreMock.apply(butler1) 

138 

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

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

141 registry1 = butler1.registry 

142 registry1.registerRun("run1") 

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

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

145 

146 # Create the collection chain 

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

148 self.assertChain( 

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

150 ) 

151 

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

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

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

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

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

157 # the command line interface. 

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

159 

160 expected = Table( 

161 array( 

162 ( 

163 ("calibration1", "CALIBRATION"), 

164 ("chain1", "CHAINED"), 

165 (" tag1", "TAGGED"), 

166 (" run1", "RUN"), 

167 (" chain2", "CHAINED"), 

168 (" calibration1", "CALIBRATION"), 

169 (" run1", "RUN"), 

170 ("chain2", "CHAINED"), 

171 (" calibration1", "CALIBRATION"), 

172 (" run1", "RUN"), 

173 ("imported_g", "RUN"), 

174 ("imported_r", "RUN"), 

175 ("run1", "RUN"), 

176 ("tag1", "TAGGED"), 

177 ) 

178 ), 

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

180 ) 

181 self.assertAstropyTablesEqual(table, expected) 

182 

183 # Test table with inverse == True 

184 table = queryCollections( 

185 "here", 

186 glob=(), 

187 collection_type=CollectionType.all(), 

188 chains="INVERSE-TREE", 

189 ) 

190 expected = Table( 

191 array( 

192 ( 

193 ("calibration1", "CALIBRATION"), 

194 (" chain2", "CHAINED"), 

195 (" chain1", "CHAINED"), 

196 ("chain1", "CHAINED"), 

197 ("chain2", "CHAINED"), 

198 (" chain1", "CHAINED"), 

199 ("imported_g", "RUN"), 

200 ("imported_r", "RUN"), 

201 ("run1", "RUN"), 

202 (" chain1", "CHAINED"), 

203 (" chain2", "CHAINED"), 

204 (" chain1", "CHAINED"), 

205 ("tag1", "TAGGED"), 

206 (" chain1", "CHAINED"), 

207 ) 

208 ), 

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

210 ) 

211 self.assertAstropyTablesEqual(table, expected) 

212 

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

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

215 expected = Table( 

216 array( 

217 ( 

218 ("calibration1", "CALIBRATION", ""), 

219 ("chain1", "CHAINED", "chain2"), 

220 ("", "", "run1"), 

221 ("", "", "tag1"), 

222 ("chain2", "CHAINED", "calibration1"), 

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

224 ("imported_g", "RUN", ""), 

225 ("imported_r", "RUN", ""), 

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

227 ("tag1", "TAGGED", ""), 

228 ) 

229 ), 

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

231 ) 

232 table = readTable(result.output) 

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

234 

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

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

237 expected = Table( 

238 array( 

239 ( 

240 ("calibration1", "CALIBRATION", "chain2"), 

241 ("chain1", "CHAINED", ""), 

242 ("chain2", "CHAINED", "chain1"), 

243 ("imported_g", "RUN", ""), 

244 ("imported_r", "RUN", ""), 

245 ("run1", "RUN", "chain1"), 

246 ("", "", "chain2"), 

247 ("tag1", "TAGGED", "chain1"), 

248 ) 

249 ), 

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

251 ) 

252 table = readTable(result.output) 

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

254 

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

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

257 expected = Table( 

258 array( 

259 ( 

260 ("calibration1", "CALIBRATION"), 

261 ("imported_g", "RUN"), 

262 ("imported_r", "RUN"), 

263 ("run1", "RUN"), 

264 ("tag1", "TAGGED"), 

265 ) 

266 ), 

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

268 ) 

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

270 

271 # Add a couple more run collections for chain testing 

272 registry1.registerRun("run2") 

273 registry1.registerRun("run3") 

274 registry1.registerRun("run4") 

275 

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

277 

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

279 

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

281 

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

283 

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

285 

286 self.assertChain( 

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

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

289 ) 

290 

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

292 

293 

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

295 unittest.main()