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 contextlib 

25from io import StringIO 

26import os 

27from typing import ( 

28 Generator, 

29 Iterable, 

30 Optional, 

31 Set, 

32 TextIO, 

33 Tuple, 

34 Union, 

35) 

36import unittest 

37import unittest.mock 

38 

39from lsst.daf.butler import ( 

40 DatasetRef, 

41 Datastore, 

42 FileDataset, 

43 Registry, 

44 YamlRepoExportBackend, 

45 YamlRepoImportBackend, 

46) 

47from lsst.daf.butler.transfers import RepoExportContext 

48from lsst.daf.butler.registry import 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 

73class TransfersTestCase(unittest.TestCase): 

74 """Tests for repository import/export functionality. 

75 """ 

76 

77 def makeRegistry(self) -> Registry: 

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

79 

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

81 """ 

82 config = RegistryConfig() 

83 config["db"] = "sqlite://" 

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

85 

86 def runImport(self, file: Union[str, TextIO], 

87 *, 

88 registry: Optional[Registry] = None, 

89 datastore: Optional[Datastore] = None, 

90 directory: Optional[str] = None, 

91 transfer: Optional[str] = None, 

92 skip_dimensions: Optional[Set[str]] = None) -> Tuple[Registry, Datastore]: 

93 """Import repository data from an export file. 

94 

95 Parameters 

96 ---------- 

97 file: `str` or `TextIO` 

98 Name of the (YAML) file that describes the data to import, or an 

99 open file-like object pointing to it. 

100 stream: ` 

101 registry: `Registry`, optional 

102 Registry instance to load datsets into. If not provided, 

103 `makeRegistry` is called. 

104 datastore: `Datastore`, optional 

105 Datastore instance to load datasets into (may be a mock). If not 

106 provided, a mock is created. 

107 directory: `str`, optional 

108 Directory containing files to import. Ignored if ``datastore`` is 

109 a mock. 

110 transfer: `str`, optional 

111 Transfer mode. See `Datastore.ingest`. 

112 skip_dimensions: `Set` [ `str` ], optional 

113 Set of dimension element names for which records should not be 

114 imported. 

115 

116 Returns 

117 ------- 

118 registry: `Registry` 

119 The `Registry` that datasets were loaded into. 

120 datastore: `Datastore` 

121 The `Datastore` instance or mock that datasets were loaded into. 

122 """ 

123 if registry is None: 

124 registry = self.makeRegistry() 

125 if datastore is None: 

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

127 if isinstance(file, str): 

128 with open(file, 'r') as stream: 

129 backend = YamlRepoImportBackend(stream, registry) 

130 backend.register() 

131 backend.load(datastore, directory=directory, transfer=transfer, 

132 skip_dimensions=skip_dimensions) 

133 else: 

134 backend = YamlRepoImportBackend(file, registry) 

135 backend.register() 

136 backend.load(datastore, directory=directory, transfer=transfer, skip_dimensions=skip_dimensions) 

137 return (registry, datastore) 

138 

139 @contextlib.contextmanager 

140 def runExport(self, *, 

141 registry: Optional[Registry] = None, 

142 datastore: Optional[Datastore] = None, 

143 stream: Optional[TextIO] = None, 

144 directory: Optional[str] = None, 

145 transfer: Optional[str] = None) -> Generator[RepoExportContext, None, None]: 

146 """Export repository data to an export file. 

147 

148 Parameters 

149 ---------- 

150 registry: `Registry`, optional 

151 Registry instance to load datasets from. If not provided, 

152 `makeRegistry` is called. 

153 datastore: `Datastore`, optional 

154 Datastore instance to load datasets from (may be a mock). If not 

155 provided, a mock is created. 

156 directory: `str`, optional 

157 Directory to contain exported file. Ignored if ``datastore`` is 

158 a mock. 

159 stream : `TextIO`, optional 

160 Writeable file-like object pointing at the export file to write. 

161 transfer: `str`, optional 

162 Transfer mode. See `Datastore.ingest`. 

163 

164 Yields 

165 ------ 

166 context : `RepoExportContext` 

167 A helper object that can be used to export repo data. This is 

168 wrapped in a context manager via the `contextlib.contextmanager` 

169 decorator. 

170 """ 

171 if registry is None: 

172 registry = self.makeRegistry() 

173 if datastore is None: 

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

175 datastore.export = _mock_export 

176 backend = YamlRepoExportBackend(stream) 

177 try: 

178 helper = RepoExportContext(registry, datastore, backend=backend, 

179 directory=directory, transfer=transfer) 

180 yield helper 

181 except BaseException: 

182 raise 

183 else: 

184 helper._finish() 

185 

186 def testReadBackwardsCompatibility(self): 

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

188 and commit to the daf_butler git repo. 

189 

190 Notes 

191 ----- 

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

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

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

195 the export format required for CALIBRATION collections to land. 

196 """ 

197 registry, _ = self.runImport(os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml")) 

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

199 # the above does not raise. 

200 self.assertGreaterEqual( 

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

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

203 ) 

204 self.assertGreaterEqual( 

205 { 

206 (record.id, record.physical_filter) 

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

208 }, 

209 { 

210 (27136, 'HSC-Z'), 

211 (11694, 'HSC-G'), 

212 (23910, 'HSC-R'), 

213 (11720, 'HSC-Y'), 

214 (23900, 'HSC-R'), 

215 (22646, 'HSC-Y'), 

216 (1248, 'HSC-I'), 

217 (19680, 'HSC-I'), 

218 (1240, 'HSC-I'), 

219 (424, 'HSC-Y'), 

220 (19658, 'HSC-I'), 

221 (344, 'HSC-Y'), 

222 (1218, 'HSC-R'), 

223 (1190, 'HSC-Z'), 

224 (23718, 'HSC-R'), 

225 (11700, 'HSC-G'), 

226 (26036, 'HSC-G'), 

227 (23872, 'HSC-R'), 

228 (1170, 'HSC-Z'), 

229 (1876, 'HSC-Y'), 

230 } 

231 ) 

232 

233 def testAllDatasetsRoundTrip(self): 

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

235 back in again. 

236 """ 

237 # Import data to play with. 

238 registry1, _ = self.runImport(os.path.join(TESTDIR, "data", "registry", "base.yaml")) 

239 self.runImport(os.path.join(TESTDIR, "data", "registry", "datasets.yaml"), registry=registry1) 

240 # Export all datasets. 

241 exportStream = StringIO() 

242 with self.runExport(stream=exportStream, registry=registry1) as exporter: 

243 exporter.saveDatasets( 

244 registry1.queryDatasets(..., collections=...) 

245 ) 

246 # Import it all again. 

247 importStream = StringIO(exportStream.getvalue()) 

248 registry2, _ = self.runImport(importStream) 

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

250 # not care about dataset_id values, which may be rewritten. 

251 self.assertCountEqual( 

252 [ref.unresolved() for ref in registry1.queryDatasets(..., collections=...)], 

253 [ref.unresolved() for ref in registry2.queryDatasets(..., collections=...)], 

254 ) 

255 

256 

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

258 unittest.main()