Hide keyboard shortcuts

Hot-keys 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

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 

25from astropy.table import Table 

26from numpy import array 

27import os 

28import unittest 

29from typing import List 

30 

31from lsst.daf.butler import ( 

32 Butler, 

33 CollectionType, 

34) 

35from lsst.daf.butler.cli.butler import cli 

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

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

38from lsst.daf.butler.script import queryCollections 

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

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

41 

42 

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

44 

45 

46class QueryCollectionsCmdTest(CliCmdTestBase, unittest.TestCase): 

47 

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

49 

50 @staticmethod 

51 def defaultExpected(): 

52 return dict(repo=None, 

53 collection_type=tuple(CollectionType.__members__.values()), 

54 chains="TABLE", 

55 glob=()) 

56 

57 @staticmethod 

58 def command(): 

59 return query_collections 

60 

61 def test_minimal(self): 

62 """Test only the required parameters, and omit the optional parameters. 

63 """ 

64 self.run_test(["query-collections", "here"], 

65 self.makeExpected(repo="here")) 

66 

67 def test_all(self): 

68 """Test all parameters""" 

69 self.run_test(["query-collections", "here", "foo*", 

70 "--collection-type", "TAGGED", 

71 "--collection-type", "RUN"], 

72 self.makeExpected(repo="here", 

73 glob=("foo*",), 

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

75 chains="TABLE")) 

76 

77 

78class QueryCollectionsScriptTest(ButlerTestHelper, unittest.TestCase): 

79 

80 def setUp(self): 

81 self.runner = LogCliRunner() 

82 

83 def testGetCollections(self): 

84 run = "ingest/run" 

85 tag = "tag" 

86 with self.runner.isolated_filesystem(): 

87 butlerCfg = Butler.makeRepo("here") 

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

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

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

91 

92 # Verify collections that were created are found by 

93 # query-collections. 

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

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

96 expected = Table((("ingest/run", "tag"), ("RUN", "TAGGED")), 

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

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

99 

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

101 # name matches with the specified pattern are returned. 

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

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

104 expected = Table((("tag",), ("TAGGED",)), 

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

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

107 

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

109 # that type are returned. 

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

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

112 expected = Table((("ingest/run",), ("RUN",)), 

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

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

115 

116 

117class ChainedCollectionsTest(ButlerTestHelper, unittest.TestCase): 

118 

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"], 

148 "[calibration1, run1]") 

149 self.assertChain(["--mode", "redefine", "chain1", "tag1", "run1", "chain2"], 

150 "[tag1, run1, chain2]") 

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 # self.assertEqual(result.exit_code, 0, clickResultMsg(result)) 

161 expected = Table(array((("imported_g", "RUN"), 

162 ("imported_r", "RUN"), 

163 ("run1", "RUN"), 

164 ("tag1", "TAGGED"), 

165 ("calibration1", "CALIBRATION"), 

166 ("chain2", "CHAINED"), 

167 (" calibration1", "CALIBRATION"), 

168 (" run1", "RUN"), 

169 ("chain1", "CHAINED"), 

170 (" tag1", "TAGGED"), 

171 (" run1", "RUN"), 

172 (" chain2", "CHAINED"), 

173 (" calibration1", "CALIBRATION"), 

174 (" run1", "RUN"))), 

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

176 self.assertAstropyTablesEqual(table, expected) 

177 

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

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

180 expected = Table(array(( 

181 ("imported_g", "RUN", ""), 

182 ("imported_r", "RUN", ""), 

183 ("run1", "RUN", ""), 

184 ("tag1", "TAGGED", ""), 

185 ("calibration1", "CALIBRATION", ""), 

186 ("chain2", "CHAINED", "[calibration1, run1]"), 

187 ("chain1", "CHAINED", "[tag1, run1, chain2]"))), 

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

189 table = readTable(result.output) 

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

191 

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

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

194 expected = Table(array(( 

195 ("imported_g", "RUN"), 

196 ("imported_r", "RUN"), 

197 ("run1", "RUN"), 

198 ("tag1", "TAGGED"), 

199 ("calibration1", "CALIBRATION"), 

200 ("calibration1", "CALIBRATION"), 

201 ("run1", "RUN"), 

202 ("tag1", "TAGGED"), 

203 ("run1", "RUN"), 

204 ("calibration1", "CALIBRATION"))), 

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

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

207 

208 # Add a couple more run collections for chain testing 

209 registry1.registerRun("run2") 

210 registry1.registerRun("run3") 

211 

212 self.assertChain(["--mode", "pop", "chain1"], 

213 "[run1, chain2]") 

214 

215 self.assertChain(["--mode", "extend", "chain1", "run2", "run3"], 

216 "[run1, chain2, run2, run3]") 

217 

218 self.assertChain(["--mode", "remove", "chain1", "chain2", "run2"], 

219 "[run1, run3]") 

220 

221 self.assertChain(["--mode", "prepend", "chain1", "chain2", "run2"], 

222 "[chain2, run2, run1, run3]") 

223 

224 self.assertChain(["--mode", "pop", "chain1", "1", "3"], 

225 "[chain2, run1]") 

226 

227 self.assertChain(["--mode", "redefine", "chain1", "chain2", "run2", "run3", "--flatten"], 

228 "[calibration1, run1, run2, run3]") 

229 

230 self.assertChain(["--mode", "pop", "chain1", "--", "-1", "-3"], 

231 "[calibration1, run2]") 

232 

233 

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

235 unittest.main()