Coverage for tests/test_transfers.py : 26%

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/>.
22from __future__ import annotations
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
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
51TESTDIR = os.path.abspath(os.path.dirname(__file__))
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.
60 This can be used to construct a `Datastore` mock that can be used in
61 repository export via::
63 datastore = unittest.mock.Mock(spec=Datastore)
64 datastore.export = _mock_export
66 """
67 for ref in refs:
68 yield FileDataset(refs=[ref],
69 path="mock/path",
70 formatter="lsst.daf.butler.formatters.json.JsonFormatter")
73class TransfersTestCase(unittest.TestCase):
74 """Tests for repository import/export functionality.
75 """
77 def makeRegistry(self) -> Registry:
78 """Create a new `Registry` instance.
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)
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.
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.
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)
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.
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`.
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()
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.
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 )
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 )
257if __name__ == "__main__": 257 ↛ 258line 257 didn't jump to line 258, because the condition on line 257 was never true
258 unittest.main()