Coverage for tests/test_registry_dataset_type_overrides.py: 100%
102 statements
« prev ^ index » next coverage.py v7.16.0, created at 2026-09-02 05:09 -0400
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 software is dual licensed under the GNU General Public License and also
10# under a 3-clause BSD license. Recipients may choose which of these licenses
11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt,
12# respectively. If you choose the GPL option then the following text applies
13# (but note that there is still no warranty even if you opt for BSD instead):
14#
15# This program is free software: you can redistribute it and/or modify
16# it under the terms of the GNU General Public License as published by
17# the Free Software Foundation, either version 3 of the License, or
18# (at your option) any later version.
19#
20# This program is distributed in the hope that it will be useful,
21# but WITHOUT ANY WARRANTY; without even the implied warranty of
22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
23# GNU General Public License for more details.
24#
25# You should have received a copy of the GNU General Public License
26# along with this program. If not, see <http://www.gnu.org/licenses/>.
28import os
29import unittest
31from lsst.daf.butler import Butler, Config, DatasetRef, DatasetType, MissingDatasetTypeError
32from lsst.daf.butler.registry import ConflictingDefinitionError
33from lsst.daf.butler.tests._repo_template_cache import make_repo_for_test
34from lsst.daf.butler.tests.dict_convertible_model import DictConvertibleModel
35from lsst.daf.butler.tests.utils import safeTestTempDir
37TESTDIR = os.path.abspath(os.path.dirname(__file__))
40class RegistryDatasetTypeOverridesTestCase(unittest.TestCase):
41 """Tests for overriding dataset type names and storage classes in registry
42 configuration.
43 """
45 def setUp(self) -> None:
46 model_type_name = "lsst.daf.butler.tests.dict_convertible_model.DictConvertibleModel"
47 config = Config()
48 # Configure two storage classes that are bidirectional convertible
49 # but with different Python types, and a third that is not
50 # convertible to either.
51 config["storageClasses"] = {
52 "BuiltinDict": {
53 "pytype": "dict",
54 "converters": {model_type_name: f"{model_type_name}.to_dict"},
55 },
56 "DictModel": {
57 "pytype": model_type_name,
58 "converters": {"dict": f"{model_type_name}.from_dict"},
59 },
60 "BuiltinList": {"pytype": "list"},
61 }
62 config["datastore"] = {
63 "formatters": {
64 "BuiltinDict": "lsst.daf.butler.formatters.yaml.YamlFormatter",
65 "DictModel": "lsst.daf.butler.formatters.json.JsonFormatter",
66 "BuiltinList": "lsst.daf.butler.formatters.json.JsonFormatter",
67 }
68 }
69 repo_dir = self.enterContext(safeTestTempDir(TESTDIR))
70 self.base_config = make_repo_for_test(root=repo_dir, config=config)
71 self.base_butler = self.enterContext(Butler.from_config(self.base_config, writeable=True))
73 def test_rename(self) -> None:
74 """Test renaming two dataset types in config overrides, to mimic the
75 case where and old dataset type is squatting on a name we want to use
76 for a new one.
77 """
78 self.setup_rename()
79 self.base_butler.registry.registerDatasetType(self.legacy_dst)
80 self.base_butler.registry.registerDatasetType(self.future_dst)
81 with Butler.from_config(self.override_config) as override_butler:
82 # Butler with overrides sees the rename:
83 self.assertEqual(override_butler.get_dataset_type("legacy_dst"), self.legacy_dst_override)
84 self.assertEqual(override_butler.get_dataset_type("dst"), self.future_dst_override)
85 with self.assertRaises(MissingDatasetTypeError):
86 override_butler.get_dataset_type("future_dst")
87 # Original butler still sees the original definitions:
88 self.assertEqual(self.base_butler.get_dataset_type("dst"), self.legacy_dst)
89 self.assertEqual(self.base_butler.get_dataset_type("future_dst"), self.future_dst)
90 with self.assertRaises(MissingDatasetTypeError):
91 self.base_butler.get_dataset_type("legacy_dst")
92 # Test dataset type queries with a new butler to hit different
93 # caching/fetching paths.
94 with Butler.from_config(self.override_config) as override_butler:
95 self.assertCountEqual(
96 override_butler.registry.queryDatasetTypes("*"),
97 [self.legacy_dst_override, self.future_dst_override],
98 )
99 # Test put/get and dataset queries across the two butlers.
100 with Butler.from_config(self.override_config, writeable=True) as override_butler:
101 legacy_ref = self.base_butler.put({"one": 1}, "dst", run="run1")
102 future_ref_override = override_butler.put([2], "dst", instrument="Cam1", run="run1")
103 self.base_butler.registry.refresh()
104 override_butler.registry.refresh()
105 self.assertEqual(self.base_butler.get("future_dst", instrument="Cam1", collections=["run1"]), [2])
106 self.assertEqual(override_butler.get("legacy_dst", collections=["run1"]), {"one": 1})
107 self.assertEqual(
108 self.base_butler.getURI("future_dst", instrument="Cam1", collections=["run1"]).getExtension(),
109 ".json",
110 )
111 self.assertEqual(
112 override_butler.getURI("legacy_dst", collections=["run1"]).getExtension(), ".yaml"
113 )
114 self.assertEqual(
115 self.base_butler.query_datasets("future_dst", collections=["run1"]),
116 [DatasetRef(self.future_dst, future_ref_override.dataId, "run1", id=future_ref_override.id)],
117 )
118 self.assertEqual(
119 override_butler.query_datasets("legacy_dst", collections=["run1"]),
120 [DatasetRef(self.legacy_dst_override, legacy_ref.dataId, "run1", id=legacy_ref.id)],
121 )
123 def test_rename_registration(self) -> None:
124 """Test dataset type registration when the dataset type name has been
125 configured to be renamed.
126 """
127 self.setup_rename()
128 with Butler.from_config(self.override_config, writeable=True) as override_butler:
129 with self.assertRaises(ConflictingDefinitionError):
130 # Even though the original name of a renamed dataset type
131 # appears missing, we shouldn't be able to register it (this
132 # is the only way in which an override rename should differ
133 # from a real DB-level rename).
134 override_butler.registry.registerDatasetType(self.future_dst)
135 # We can register the original dataset type in the base repo by
136 # registering its rename in the override repo.
137 override_butler.registry.registerDatasetType(self.future_dst_override)
138 self.base_butler.registry.refresh()
139 self.assertEqual(self.base_butler.get_dataset_type("future_dst"), self.future_dst)
141 def setup_rename(self) -> None:
142 """Do additional test setup for rename-override tests."""
143 self.legacy_dst = DatasetType(
144 "dst", dimensions=(), storageClass="BuiltinDict", universe=self.base_butler.dimensions
145 )
146 self.future_dst = DatasetType(
147 "future_dst",
148 dimensions={"instrument"},
149 storageClass="BuiltinList",
150 universe=self.base_butler.dimensions,
151 )
152 self.base_butler.registry.insertDimensionData("instrument", {"name": "Cam1"})
153 self.base_butler.collections.register("run1")
154 self.override_config = self.base_config.copy()
155 self.override_config[".registry.managers.datasets"] = {
156 "cls": self.base_config[".registry.managers.datasets"],
157 "config": {
158 "overrides": {
159 "dst": {"rename": "legacy_dst"},
160 "future_dst": {"rename": "dst"},
161 }
162 },
163 }
164 self.legacy_dst_override = DatasetType(
165 "legacy_dst", dimensions=(), storageClass="BuiltinDict", universe=self.base_butler.dimensions
166 )
167 self.future_dst_override = DatasetType(
168 "dst",
169 dimensions={"instrument"},
170 storageClass="BuiltinList",
171 universe=self.base_butler.dimensions,
172 )
174 def test_storage_class_override_config(self) -> None:
175 """Test overriding the storage class of a dataset type via a repository
176 configuration override.
177 """
178 self.setup_storage_class()
179 self.base_butler.registry.registerDatasetType(self.base_dst)
180 with Butler.from_config(self.override_config, writeable=True) as override_butler:
181 self.assertEqual(self.base_butler.get_dataset_type("dst"), self.base_dst)
182 self.assertEqual(override_butler.get_dataset_type("dst"), self.override_dst)
183 # 'put' with both butlers using their natural storage classes.
184 obj1 = {"one": "1"}
185 obj2 = DictConvertibleModel(content={"two": "2"}, extra="three")
186 ref1 = self.base_butler.put(obj1, "dst", run="run1")
187 ref2 = override_butler.put(obj2, "dst", run="run2")
188 self.base_butler.registry.refresh()
189 override_butler.registry.refresh()
190 # We can 'get' using the refs with either butler, since that's
191 # just a call-level dataset type override that makes the datastore
192 # 'get' match the storage class used for the write.
193 self.assertEqual(self.base_butler.get(ref1), obj1)
194 self.assertEqual(self.base_butler.get(ref2), obj2)
195 self.assertEqual(override_butler.get(ref1), obj1)
196 self.assertEqual(override_butler.get(ref2), obj2)
197 # Check that we used the right formatters by looking at the
198 # extensions.
199 self.assertEqual(self.base_butler.getURI(ref1).getExtension(), ".yaml")
200 self.assertEqual(self.base_butler.getURI(ref2).getExtension(), ".json")
201 # Do a 'get' with just the dataset type name with the same butler
202 # we used to do the write; no conversion necessary.
203 self.assertEqual(self.base_butler.get("dst", collections=["run1"]), obj1)
204 self.assertEqual(override_butler.get("dst", collections=["run2"]), obj2)
205 # Do a 'get' with just the dataset type with the other butlers.
206 # This should do a storage class conversion.
207 self.assertEqual(self.base_butler.get("dst", collections=["run2"]), obj2.to_dict())
208 self.assertEqual(
209 override_butler.get("dst", collections=["run1"]), DictConvertibleModel.from_dict(obj1)
210 )
211 # Do a 'get' with call-level storage class overrides that differ
212 # from the client's storage class.
213 self.assertEqual(
214 self.base_butler.get("dst", collections=["run1"], storageClass="DictModel"),
215 DictConvertibleModel.from_dict(obj1),
216 )
217 self.assertEqual(
218 override_butler.get("dst", collections=["run2"], storageClass="BuiltinDict"), obj2.to_dict()
219 )
220 # Query for datasets and check that the refs have the right
221 # storage classes.
222 self.assertCountEqual(
223 self.base_butler.query_datasets("dst", collections=["run1", "run2"], find_first=False),
224 [ref1, ref2.overrideStorageClass("BuiltinDict")],
225 )
226 self.assertCountEqual(
227 override_butler.query_datasets("dst", collections=["run1", "run2"], find_first=False),
228 [ref1.overrideStorageClass("DictModel"), ref2],
229 )
231 def test_storage_class_override_config_registration(self) -> None:
232 """Test registering dataset types whose storage classes have been
233 overridden via repository configuration.
234 """
235 self.setup_storage_class()
236 with Butler.from_config(self.override_config, writeable=True) as override_butler:
237 # Trying to register a dataset type with a storage class override
238 # when it doesn't already exist is an error, because we don't know
239 # what it's storage class should really be in the database.
240 with self.assertRaises(ConflictingDefinitionError):
241 override_butler.registry.registerDatasetType(self.override_dst)
242 # Registering in the base repository is of course fine.
243 self.base_butler.registry.registerDatasetType(self.base_dst)
244 # Re-registering with the new storage class in the override repo
245 # should be a no-op, because it should look like it already exists
246 # with that storage class.
247 self.assertFalse(override_butler.registry.registerDatasetType(self.override_dst))
248 # Re-registering with the base repo's storage class in the
249 # override repo should be an error, because to the override
250 # repo it looks like it has a different storage class.
251 with self.assertRaises(ConflictingDefinitionError):
252 override_butler.registry.registerDatasetType(self.base_dst)
254 def setup_storage_class(self) -> None:
255 """Do additional test setup for storage-class-override tests."""
256 self.base_dst = DatasetType(
257 "dst", dimensions=(), storageClass="BuiltinDict", universe=self.base_butler.dimensions
258 )
259 self.override_dst = self.base_dst.overrideStorageClass("DictModel")
260 self.base_butler.collections.register("run1")
261 self.base_butler.collections.register("run2")
262 self.override_config = self.base_config.copy()
263 self.override_config[".registry.managers.datasets"] = {
264 "cls": self.base_config[".registry.managers.datasets"],
265 "config": {
266 "overrides": {
267 "dst": {"storageClass": "DictModel"},
268 }
269 },
270 }
273if __name__ == "__main__":
274 unittest.main()