Coverage for tests/test_simpleButler.py : 16%

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 Any
27import unittest
29try:
30 import numpy as np
31except ImportError:
32 np = None
34import astropy.time
36from lsst.daf.butler import (
37 Butler,
38 ButlerConfig,
39 CollectionType,
40 Registry,
41 Timespan,
42)
43from lsst.daf.butler.registry import RegistryConfig
44from lsst.daf.butler.tests import DatastoreMock
45from lsst.daf.butler.tests.utils import makeTestTempDir, removeTestTempDir
48TESTDIR = os.path.abspath(os.path.dirname(__file__))
51class SimpleButlerTestCase(unittest.TestCase):
52 """Tests for butler (including import/export functionality) that should not
53 depend on the Registry Database backend or Datastore implementation, and
54 can instead utilize an in-memory SQLite Registry and a mocked Datastore.
55 """
57 def setUp(self):
58 self.root = makeTestTempDir(TESTDIR)
60 def tearDown(self):
61 removeTestTempDir(self.root)
63 def makeButler(self, **kwargs: Any) -> Butler:
64 """Return new Butler instance on each call.
65 """
66 config = ButlerConfig()
68 # make separate temporary directory for registry of this instance
69 tmpdir = tempfile.mkdtemp(dir=self.root)
70 config["registry", "db"] = f"sqlite:///{tmpdir}/gen3.sqlite3"
71 config["root"] = self.root
73 # have to make a registry first
74 registryConfig = RegistryConfig(config.get("registry"))
75 Registry.createFromConfig(registryConfig)
77 butler = Butler(config, **kwargs)
78 DatastoreMock.apply(butler)
79 return butler
81 def testReadBackwardsCompatibility(self):
82 """Test that we can read an export file written by a previous version
83 and commit to the daf_butler git repo.
85 Notes
86 -----
87 At present this export file includes only dimension data, not datasets,
88 which greatly limits the usefulness of this test. We should address
89 this at some point, but I think it's best to wait for the changes to
90 the export format required for CALIBRATION collections to land.
91 """
92 butler = self.makeButler(writeable=True)
93 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "hsc-rc2-subset.yaml"))
94 # Spot-check a few things, but the most important test is just that
95 # the above does not raise.
96 self.assertGreaterEqual(
97 set(record.id for record in butler.registry.queryDimensionRecords("detector", instrument="HSC")),
98 set(range(104)), # should have all science CCDs; may have some focus ones.
99 )
100 self.assertGreaterEqual(
101 {
102 (record.id, record.physical_filter)
103 for record in butler.registry.queryDimensionRecords("visit", instrument="HSC")
104 },
105 {
106 (27136, 'HSC-Z'),
107 (11694, 'HSC-G'),
108 (23910, 'HSC-R'),
109 (11720, 'HSC-Y'),
110 (23900, 'HSC-R'),
111 (22646, 'HSC-Y'),
112 (1248, 'HSC-I'),
113 (19680, 'HSC-I'),
114 (1240, 'HSC-I'),
115 (424, 'HSC-Y'),
116 (19658, 'HSC-I'),
117 (344, 'HSC-Y'),
118 (1218, 'HSC-R'),
119 (1190, 'HSC-Z'),
120 (23718, 'HSC-R'),
121 (11700, 'HSC-G'),
122 (26036, 'HSC-G'),
123 (23872, 'HSC-R'),
124 (1170, 'HSC-Z'),
125 (1876, 'HSC-Y'),
126 }
127 )
129 def testDatasetTransfers(self):
130 """Test exporting all datasets from a repo and then importing them all
131 back in again.
132 """
133 # Import data to play with.
134 butler1 = self.makeButler(writeable=True)
135 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml"))
136 butler1.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml"))
137 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file:
138 # Export all datasets.
139 with butler1.export(filename=file.name) as exporter:
140 exporter.saveDatasets(
141 butler1.registry.queryDatasets(..., collections=...)
142 )
143 # Import it all again.
144 butler2 = self.makeButler(writeable=True)
145 butler2.import_(filename=file.name)
146 # Check that it all round-tripped. Use unresolved() to make
147 # comparison not care about dataset_id values, which may be
148 # rewritten.
149 self.assertCountEqual(
150 [ref.unresolved() for ref in butler1.registry.queryDatasets(..., collections=...)],
151 [ref.unresolved() for ref in butler2.registry.queryDatasets(..., collections=...)],
152 )
154 def testCollectionTransfers(self):
155 """Test exporting and then importing collections of various types.
156 """
157 # Populate a registry with some datasets.
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 registry1 = butler1.registry
162 # Add some more collections.
163 registry1.registerRun("run1")
164 registry1.registerCollection("tag1", CollectionType.TAGGED)
165 registry1.registerCollection("calibration1", CollectionType.CALIBRATION)
166 registry1.registerCollection("chain1", CollectionType.CHAINED)
167 registry1.registerCollection("chain2", CollectionType.CHAINED)
168 registry1.setCollectionChain("chain1", ["tag1", "run1", "chain2"])
169 registry1.setCollectionChain("chain2", ["calibration1", "run1"])
170 # Associate some datasets into the TAGGED and CALIBRATION collections.
171 flats1 = list(registry1.queryDatasets("flat", collections=...))
172 registry1.associate("tag1", flats1)
173 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai")
174 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai")
175 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai")
176 bias2a = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g")
177 bias3a = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g")
178 bias2b = registry1.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r")
179 bias3b = registry1.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r")
180 registry1.certify("calibration1", [bias2a, bias3a], Timespan(t1, t2))
181 registry1.certify("calibration1", [bias2b], Timespan(t2, None))
182 registry1.certify("calibration1", [bias3b], Timespan(t2, t3))
184 with tempfile.NamedTemporaryFile(mode='w', suffix=".yaml") as file:
185 # Export all collections, and some datasets.
186 with butler1.export(filename=file.name) as exporter:
187 # Sort results to put chain1 before chain2, which is
188 # intentionally not topological order.
189 for collection in sorted(registry1.queryCollections()):
190 exporter.saveCollection(collection)
191 exporter.saveDatasets(flats1)
192 exporter.saveDatasets([bias2a, bias2b, bias3a, bias3b])
193 # Import them into a new registry.
194 butler2 = self.makeButler(writeable=True)
195 butler2.import_(filename=file.name)
196 registry2 = butler2.registry
197 # Check that it all round-tripped, starting with the collections
198 # themselves.
199 self.assertIs(registry2.getCollectionType("run1"), CollectionType.RUN)
200 self.assertIs(registry2.getCollectionType("tag1"), CollectionType.TAGGED)
201 self.assertIs(registry2.getCollectionType("calibration1"), CollectionType.CALIBRATION)
202 self.assertIs(registry2.getCollectionType("chain1"), CollectionType.CHAINED)
203 self.assertIs(registry2.getCollectionType("chain2"), CollectionType.CHAINED)
204 self.assertEqual(
205 list(registry2.getCollectionChain("chain1")),
206 ["tag1", "run1", "chain2"],
207 )
208 self.assertEqual(
209 list(registry2.getCollectionChain("chain2")),
210 ["calibration1", "run1"],
211 )
212 # Check that tag collection contents are the same.
213 self.maxDiff = None
214 self.assertCountEqual(
215 [ref.unresolved() for ref in registry1.queryDatasets(..., collections="tag1")],
216 [ref.unresolved() for ref in registry2.queryDatasets(..., collections="tag1")],
217 )
218 # Check that calibration collection contents are the same.
219 self.assertCountEqual(
220 [(assoc.ref.unresolved(), assoc.timespan)
221 for assoc in registry1.queryDatasetAssociations("bias", collections="calibration1")],
222 [(assoc.ref.unresolved(), assoc.timespan)
223 for assoc in registry2.queryDatasetAssociations("bias", collections="calibration1")],
224 )
226 def testButlerGet(self):
227 """Test that butler.get can work with different variants."""
229 # Import data to play with.
230 butler = self.makeButler(writeable=True)
231 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml"))
232 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml"))
234 # Find the DatasetRef for a flat
235 coll = "imported_g"
236 flat2g = butler.registry.findDataset("flat", instrument="Cam1", detector=2, physical_filter="Cam1-G",
237 collections=coll)
239 # Create a numpy integer to check that works fine
240 detector_np = np.int64(2) if np else 2
241 print(type(detector_np))
243 # Try to get it using different variations of dataId + keyword
244 # arguments
245 # Note that instrument.class_name does not work
246 variants = (
247 (None, {"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}),
248 (None, {"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}),
249 ({"instrument": "Cam1", "detector": 2, "physical_filter": "Cam1-G"}, {}),
250 ({"instrument": "Cam1", "detector": detector_np, "physical_filter": "Cam1-G"}, {}),
251 ({"instrument": "Cam1", "detector": 2}, {"physical_filter": "Cam1-G"}),
252 ({"detector.full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}),
253 ({"full_name": "Ab"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}),
254 (None, {"full_name": "Ab", "instrument": "Cam1", "physical_filter": "Cam1-G"}),
255 ({"name_in_raft": "b", "raft": "A"}, {"instrument": "Cam1", "physical_filter": "Cam1-G"}),
256 ({"name_in_raft": "b"}, {"raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}),
257 (None, {"name_in_raft": "b", "raft": "A", "instrument": "Cam1", "physical_filter": "Cam1-G"}),
258 ({"detector.name_in_raft": "b", "detector.raft": "A"},
259 {"instrument": "Cam1", "physical_filter": "Cam1-G"}),
260 ({"detector.name_in_raft": "b", "detector.raft": "A",
261 "instrument": "Cam1", "physical_filter": "Cam1-G"}, {}),
262 )
264 for dataId, kwds in variants:
265 try:
266 flat_id, _ = butler.get("flat", dataId=dataId, collections=coll, **kwds)
267 except Exception as e:
268 raise type(e)(f"{str(e)}: dataId={dataId}, kwds={kwds}") from e
269 self.assertEqual(flat_id, flat2g.id, msg=f"DataId: {dataId}, kwds: {kwds}")
271 def testGetCalibration(self):
272 """Test that `Butler.get` can be used to fetch from
273 `~CollectionType.CALIBRATION` collections if the data ID includes
274 extra dimensions with temporal information.
275 """
276 # Import data to play with.
277 butler = self.makeButler(writeable=True)
278 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "base.yaml"))
279 butler.import_(filename=os.path.join(TESTDIR, "data", "registry", "datasets.yaml"))
280 # Certify some biases into a CALIBRATION collection.
281 registry = butler.registry
282 registry.registerCollection("calibs", CollectionType.CALIBRATION)
283 t1 = astropy.time.Time('2020-01-01T01:00:00', format="isot", scale="tai")
284 t2 = astropy.time.Time('2020-01-01T02:00:00', format="isot", scale="tai")
285 t3 = astropy.time.Time('2020-01-01T03:00:00', format="isot", scale="tai")
286 bias2a = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_g")
287 bias3a = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_g")
288 bias2b = registry.findDataset("bias", instrument="Cam1", detector=2, collections="imported_r")
289 bias3b = registry.findDataset("bias", instrument="Cam1", detector=3, collections="imported_r")
290 registry.certify("calibs", [bias2a, bias3a], Timespan(t1, t2))
291 registry.certify("calibs", [bias2b], Timespan(t2, None))
292 registry.certify("calibs", [bias3b], Timespan(t2, t3))
293 # Insert some exposure dimension data.
294 registry.insertDimensionData(
295 "exposure",
296 {
297 "instrument": "Cam1",
298 "id": 3,
299 "obs_id": "three",
300 "timespan": Timespan(t1, t2),
301 "physical_filter": "Cam1-G",
302 "day_obs": 20201114,
303 "seq_num": 55,
304 },
305 {
306 "instrument": "Cam1",
307 "id": 4,
308 "obs_id": "four",
309 "timespan": Timespan(t2, t3),
310 "physical_filter": "Cam1-G",
311 "day_obs": 20211114,
312 "seq_num": 42,
313 },
314 )
315 # Get some biases from raw-like data IDs.
316 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 3, "detector": 2},
317 collections="calibs")
318 self.assertEqual(bias2a_id, bias2a.id)
319 bias3b_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure": 4, "detector": 3},
320 collections="calibs")
321 self.assertEqual(bias3b_id, bias3b.id)
323 # Get using the kwarg form
324 bias3b_id, _ = butler.get("bias",
325 instrument="Cam1", exposure=4, detector=3,
326 collections="calibs")
327 self.assertEqual(bias3b_id, bias3b.id)
329 # Do it again but using the record information
330 bias2a_id, _ = butler.get("bias", {"instrument": "Cam1", "exposure.obs_id": "three",
331 "detector.full_name": "Ab"},
332 collections="calibs")
333 self.assertEqual(bias2a_id, bias2a.id)
334 bias3b_id, _ = butler.get("bias", {"exposure.obs_id": "four",
335 "detector.full_name": "Ba"},
336 collections="calibs", instrument="Cam1")
337 self.assertEqual(bias3b_id, bias3b.id)
339 # And again but this time using the alternate value rather than
340 # the primary.
341 bias3b_id, _ = butler.get("bias", {"exposure": "four",
342 "detector": "Ba"},
343 collections="calibs", instrument="Cam1")
344 self.assertEqual(bias3b_id, bias3b.id)
346 # And again but this time using the alternate value rather than
347 # the primary and do it in the keyword arguments.
348 bias3b_id, _ = butler.get("bias",
349 exposure="four", detector="Ba",
350 collections="calibs", instrument="Cam1")
351 self.assertEqual(bias3b_id, bias3b.id)
353 # Now with implied record columns
354 bias3b_id, _ = butler.get("bias", day_obs=20211114, seq_num=42,
355 raft="B", name_in_raft="a",
356 collections="calibs", instrument="Cam1")
357 self.assertEqual(bias3b_id, bias3b.id)
360if __name__ == "__main__": 360 ↛ 361line 360 didn't jump to line 361, because the condition on line 360 was never true
361 unittest.main()