Coverage for tests/test_simpleButler.py : 17%

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 os
25import tempfile
26from typing import (
27 Any,
28 Iterable,
29 Mapping,
30 Optional,
31 Tuple,
32)
33import unittest
34import unittest.mock
36import astropy.time
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
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")
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)
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 """
87 def makeRegistry(self) -> Registry:
88 """Create a new `Registry` instance.
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)
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
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.
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 )
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 )
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))
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 )
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)
299if __name__ == "__main__": 299 ↛ 300line 299 didn't jump to line 300, because the condition on line 299 was never true
300 unittest.main()