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 

22from __future__ import annotations 

23 

24import os 

25import tempfile 

26from typing import ( 

27 Any, 

28 Iterable, 

29 Mapping, 

30 Optional, 

31 Tuple, 

32) 

33import unittest 

34import unittest.mock 

35 

36import astropy.time 

37 

38from lsst.daf.butler import ( 

39 Butler, 

40 ButlerConfig, 

41 CollectionType, 

42 DatasetRef, 

43 Datastore, 

44 FileDataset, 

45 Registry, 

46 Timespan, 

47) 

48from lsst.daf.butler.registry import CollectionSearch, RegistryConfig 

49 

50 

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

52 

53 

54def _mock_export(refs: Iterable[DatasetRef], *, 

55 directory: Optional[str] = None, 

56 transfer: Optional[str] = None) -> Iterable[FileDataset]: 

57 """A mock of `Datastore.export` that satisifies the requirement that the 

58 refs passed in are included in the `FileDataset` objects returned. 

59 

60 This can be used to construct a `Datastore` mock that can be used in 

61 repository export via:: 

62 

63 datastore = unittest.mock.Mock(spec=Datastore) 

64 datastore.export = _mock_export 

65 

66 """ 

67 for ref in refs: 

68 yield FileDataset(refs=[ref], 

69 path="mock/path", 

70 formatter="lsst.daf.butler.formatters.json.JsonFormatter") 

71 

72 

73def _mock_get(ref: DatasetRef, parameters: Optional[Mapping[str, Any]] = None 

74 ) -> Tuple[int, Optional[Mapping[str, Any]]]: 

75 """A mock of `Datastore.get` that just returns the integer dataset ID value 

76 and parameters it was given. 

77 """ 

78 return (ref.id, parameters) 

79 

80 

81class SimpleButlerTestCase(unittest.TestCase): 

82 """Tests for butler (including import/export functionality) that should not 

83 depend on the Registry Database backend or Datastore implementation, and 

84 can instead utilize an in-memory SQLite Registry and a mocked Datastore. 

85 """ 

86 

87 def makeRegistry(self) -> Registry: 

88 """Create a new `Registry` instance. 

89 

90 The default implementation returns a SQLite in-memory database. 

91 """ 

92 config = RegistryConfig() 

93 config["db"] = "sqlite:///:memory:" 

94 return Registry.fromConfig(config, create=True) 

95 

96 def makeButler(self, **kwargs: Any) -> Butler: 

97 config = ButlerConfig() 

98 config["registry", "db"] = "sqlite:///:memory:" 

99 with unittest.mock.patch.object(Datastore, "fromConfig", spec=Datastore.fromConfig): 

100 butler = Butler(config, **kwargs) 

101 butler.datastore.export = _mock_export 

102 butler.datastore.get = _mock_get 

103 return butler 

104 

105 def testReadBackwardsCompatibility(self): 

106 """Test that we can read an export file written by a previous version 

107 and commit to the daf_butler git repo. 

108 

109 Notes 

110 ----- 

111 At present this export file includes only dimension data, not datasets, 

112 which greatly limits the usefulness of this test. We should address 

113 this at some point, but I think it's best to wait for the changes to 

114 the export format required for CALIBRATION collections to land. 

115 """ 

116 butler = self.makeButler(writeable=True) 

117 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

118 # Spot-check a few things, but the most important test is just that 

119 # the above does not raise. 

120 self.assertGreaterEqual( 

121 set(record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")), 

122 set(range(104)), # should have all science CCDs; may have some focus ones. 

123 ) 

124 self.assertGreaterEqual( 

125 { 

126 (record.id, record.physical_filter) 

127 for record in butler.registry.queryDimensionRecords("visit", instrument="HSC") 

128 }, 

129 { 

130 (27136, 'HSC-Z'), 

131 (11694, 'HSC-G'), 

132 (23910, 'HSC-R'), 

133 (11720, 'HSC-Y'), 

134 (23900, 'HSC-R'), 

135 (22646, 'HSC-Y'), 

136 (1248, 'HSC-I'), 

137 (19680, 'HSC-I'), 

138 (1240, 'HSC-I'), 

139 (424, 'HSC-Y'), 

140 (19658, 'HSC-I'), 

141 (344, 'HSC-Y'), 

142 (1218, 'HSC-R'), 

143 (1190, 'HSC-Z'), 

144 (23718, 'HSC-R'), 

145 (11700, 'HSC-G'), 

146 (26036, 'HSC-G'), 

147 (23872, 'HSC-R'), 

148 (1170, 'HSC-Z'), 

149 (1876, 'HSC-Y'), 

150 } 

151 ) 

152 

153 def testDatasetTransfers(self): 

154 """Test exporting all datasets from a repo and then importing them all 

155 back in again. 

156 """ 

157 # Import data to play with. 

158 butler1 = self.makeButler(writeable=True) 

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

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

161 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

162 # Export all datasets. 

163 with butler1.export(filename=file.name) as exporter: 

164 exporter.saveDatasets( 

165 butler1.registry.queryDatasets(..., collections=...) 

166 ) 

167 # Import it all again. 

168 butler2 = self.makeButler(writeable=True) 

169 butler2.import_(filename=file.name) 

170 # Check that it all round-tripped. Use unresolved() to make 

171 # comparison not care about dataset_id values, which may be 

172 # rewritten. 

173 self.assertCountEqual( 

174 [ref.unresolved() for ref in butler1.registry.queryDatasets(..., collections=...)], 

175 [ref.unresolved() for ref in butler2.registry.queryDatasets(..., collections=...)], 

176 ) 

177 

178 def testCollectionTransfers(self): 

179 """Test exporting and then importing collections of various types. 

180 """ 

181 # Populate a registry with some datasets. 

182 butler1 = self.makeButler(writeable=True) 

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

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

185 registry1 = butler1.registry 

186 # Add some more collections. 

187 registry1.registerRun("run1") 

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

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

190 registry1.registerCollection("chain1", CollectionType.CHAINED) 

191 registry1.registerCollection("chain2", CollectionType.CHAINED) 

192 registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"]) 

193 registry1.setCollectionChain("chain2", [("calibration1", ["bias"]), "run1"]) 

194 # Associate some datasets into the TAGGED and CALIBRATION collections. 

195 flats1 = list(registry1.queryDatasets("flat", collections=...)) 

196 registry1.associate("tag1", flats1) 

197 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

198 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

199 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

200 bias2a = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

201 bias3a = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

202 bias2b = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

203 bias3b = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

204 registry1.certify("calibration1", [bias2a, bias3a], Timespan(t1, t2)) 

205 registry1.certify("calibration1", [bias2b], Timespan(t2, None)) 

206 registry1.certify("calibration1", [bias3b], Timespan(t2, t3)) 

207 

208 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file: 

209 # Export all collections, and some datasets. 

210 with butler1.export(filename=file.name) as exporter: 

211 # Sort results to put chain1 before chain2, which is 

212 # intentionally not topological order. 

213 for collection in sorted(registry1.queryCollections()): 

214 exporter.saveCollection(collection) 

215 exporter.saveDatasets(flats1) 

216 exporter.saveDatasets([bias2a, bias2b, bias3a, bias3b]) 

217 # Import them into a new registry. 

218 butler2 = self.makeButler(writeable=True) 

219 butler2.import_(filename=file.name) 

220 registry2 = butler2.registry 

221 # Check that it all round-tripped, starting with the collections 

222 # themselves. 

223 self.assertIs(registry2.getCollectionType("run1"), CollectionType.RUN) 

224 self.assertIs(registry2.getCollectionType("tag1"), CollectionType.TAGGED) 

225 self.assertIs(registry2.getCollectionType("calibration1"), CollectionType.CALIBRATION) 

226 self.assertIs(registry2.getCollectionType("chain1"), CollectionType.CHAINED) 

227 self.assertIs(registry2.getCollectionType("chain2"), CollectionType.CHAINED) 

228 self.assertEqual( 

229 registry2.getCollectionChain("chain1"), 

230 CollectionSearch.fromExpression(["tag1", "run1", "chain2"]), 

231 ) 

232 self.assertEqual( 

233 registry2.getCollectionChain("chain2"), 

234 CollectionSearch.fromExpression([("calibration1", ["bias"]), "run1"]), 

235 ) 

236 # Check that tag collection contents are the same. 

237 self.maxDiff = None 

238 self.assertCountEqual( 

239 [ref.unresolved() for ref in registry1.queryDatasets(..., collections="tag1")], 

240 [ref.unresolved() for ref in registry2.queryDatasets(..., collections="tag1")], 

241 ) 

242 # Check that calibration collection contents are the same. 

243 self.assertCountEqual( 

244 [(assoc.ref.unresolved(), assoc.timespan) 

245 for assoc in registry1.queryDatasetAssociations("bias", collections="calibration1")], 

246 [(assoc.ref.unresolved(), assoc.timespan) 

247 for assoc in registry2.queryDatasetAssociations("bias", collections="calibration1")], 

248 ) 

249 

250 def testGetCalibration(self): 

251 """Test that `Butler.get` can be used to fetch from 

252 `~CollectionType.CALIBRATION` collections if the data ID includes 

253 extra dimensions with temporal information. 

254 """ 

255 # Import data to play with. 

256 butler = self.makeButler(writeable=True) 

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

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

259 # Certify some biases into a CALIBRATION collection. 

260 registry = butler.registry 

261 registry.registerCollection("calibs", CollectionType.CALIBRATION) 

262 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai") 

263 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai") 

264 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai") 

265 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g") 

266 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g") 

267 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r") 

268 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r") 

269 registry.certify("calibs", [bias2a, bias3a], Timespan(t1, t2)) 

270 registry.certify("calibs", [bias2b], Timespan(t2, None)) 

271 registry.certify("calibs", [bias3b], Timespan(t2, t3)) 

272 # Insert some exposure dimension data. 

273 registry.insertDimensionData( 

274 "exposure", 

275 { 

276 "instrument": "Cam1", 

277 "id": 3, 

278 "name": "three", 

279 "timespan": Timespan(t1, t2), 

280 "physical_filter": "Cam1-G", 

281 }, 

282 { 

283 "instrument": "Cam1", 

284 "id": 4, 

285 "name": "four", 

286 "timespan": Timespan(t2, t3), 

287 "physical_filter": "Cam1-G", 

288 }, 

289 ) 

290 # Get some biases from raw-like data IDs. 

291 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 3, "detector": 2}, 

292 collections="calibs") 

293 self.assertEqual(bias2a_id, bias2a.id) 

294 bias3b_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 4, "detector": 3}, 

295 collections="calibs") 

296 self.assertEqual(bias3b_id, bias3b.id) 

297 

298 

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

300 unittest.main()