Coverage for tests / test_obscore.py: 15%
293 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-24 01:10 -0800
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 tempfile
30import unittest
31from abc import abstractmethod
32from typing import cast
34import astropy.time
35import sqlalchemy
37from lsst.daf.butler import (
38 CollectionType,
39 Config,
40 DataCoordinate,
41 DatasetRef,
42 DatasetType,
43 StorageClassFactory,
44)
45from lsst.daf.butler.registry import RegistryConfig, _RegistryFactory
46from lsst.daf.butler.registry.obscore import (
47 DatasetTypeConfig,
48 ObsCoreConfig,
49 ObsCoreLiveTableManager,
50 ObsCoreSchema,
51 RegionTypeWarning,
52)
53from lsst.daf.butler.registry.obscore._schema import _STATIC_COLUMNS
54from lsst.daf.butler.registry.sql_registry import SqlRegistry
55from lsst.daf.butler.tests.postgresql import setup_postgres_test_db
56from lsst.daf.butler.tests.utils import TestCaseMixin, makeTestTempDir, removeTestTempDir
57from lsst.sphgeom import Box, ConvexPolygon, LonLat, UnitVector3d
59TESTDIR = os.path.abspath(os.path.dirname(__file__))
62class ObsCoreTests(TestCaseMixin):
63 """Base class for testing obscore manager functionality."""
65 root: str
67 def make_registry(
68 self, collections: list[str] | None = None, collection_type: str | None = None
69 ) -> SqlRegistry:
70 """Create new empty Registry."""
71 config = self.make_registry_config(collections, collection_type)
72 registry = _RegistryFactory(config).create_from_config(butlerRoot=self.root)
73 self.addCleanup(registry.close)
74 self.initialize_registry(registry)
75 return registry
77 @abstractmethod
78 def make_registry_config(
79 self, collections: list[str] | None = None, collection_type: str | None = None
80 ) -> RegistryConfig:
81 """Make Registry configuration."""
82 raise NotImplementedError()
84 def initialize_registry(self, registry: SqlRegistry) -> None:
85 """Populate Registry with the things that we need for tests."""
86 registry.insertDimensionData("instrument", {"name": "DummyCam"})
87 registry.insertDimensionData(
88 "physical_filter", {"instrument": "DummyCam", "name": "d-r", "band": "r"}
89 )
90 registry.insertDimensionData("day_obs", {"instrument": "DummyCam", "id": 20200101})
91 for detector in (1, 2, 3, 4):
92 registry.insertDimensionData(
93 "detector", {"instrument": "DummyCam", "id": detector, "full_name": f"detector{detector}"}
94 )
96 for exposure in (1, 2, 3, 4):
97 registry.insertDimensionData("group", {"instrument": "DummyCam", "name": f"group{exposure}"})
98 registry.insertDimensionData(
99 "exposure",
100 {
101 "instrument": "DummyCam",
102 "id": exposure,
103 "obs_id": f"exposure{exposure}",
104 "physical_filter": "d-r",
105 "group": f"group{exposure}",
106 "day_obs": 20200101,
107 },
108 )
110 registry.insertDimensionData("visit_system", {"instrument": "DummyCam", "id": 1, "name": "default"})
112 for visit in (1, 2, 3, 4, 9):
113 visit_start = astropy.time.Time(f"2020-01-01 08:0{visit}:00", scale="tai")
114 visit_end = astropy.time.Time(f"2020-01-01 08:0{visit}:45", scale="tai")
115 registry.insertDimensionData(
116 "visit",
117 {
118 "instrument": "DummyCam",
119 "id": visit,
120 "name": f"visit{visit}",
121 "physical_filter": "d-r",
122 "datetime_begin": visit_start,
123 "datetime_end": visit_end,
124 "day_obs": 20200101,
125 },
126 )
127 registry.insertDimensionData(
128 "visit_system_membership",
129 {"instrument": "DummyCam", "visit": visit, "visit_system": 1},
130 )
132 # Only couple of exposures are linked to visits.
133 for visit in (1, 2):
134 registry.insertDimensionData(
135 "visit_definition",
136 {
137 "instrument": "DummyCam",
138 "exposure": visit,
139 "visit": visit,
140 },
141 )
143 # map visit and detector to region
144 self.regions: dict[tuple[int, int], ConvexPolygon] = {}
145 for visit in (1, 2, 3, 4):
146 for detector in (1, 2, 3, 4):
147 lon = visit * 90 - 88
148 lat = detector * 2 - 5
149 region = ConvexPolygon(
150 [
151 UnitVector3d(LonLat.fromDegrees(lon - 1.0, lat - 1.0)),
152 UnitVector3d(LonLat.fromDegrees(lon + 1.0, lat - 1.0)),
153 UnitVector3d(LonLat.fromDegrees(lon + 1.0, lat + 1.0)),
154 UnitVector3d(LonLat.fromDegrees(lon - 1.0, lat + 1.0)),
155 ]
156 )
157 registry.insertDimensionData(
158 "visit_detector_region",
159 {
160 "instrument": "DummyCam",
161 "visit": visit,
162 "detector": detector,
163 "region": region,
164 },
165 )
166 self.regions[(visit, detector)] = region
168 # Visit 9 has non-polygon region
169 for detector in (1, 2, 3, 4):
170 lat = detector * 2 - 5
171 region = Box.fromDegrees(17.0, lat - 1.0, 19.0, lat + 1.0)
172 registry.insertDimensionData(
173 "visit_detector_region",
174 {
175 "instrument": "DummyCam",
176 "visit": 9,
177 "detector": detector,
178 "region": region,
179 },
180 )
182 # Add few dataset types
183 storage_class_factory = StorageClassFactory()
184 storage_class = storage_class_factory.getStorageClass("StructuredDataDict")
186 self.dataset_types: dict[str, DatasetType] = {}
188 dimensions = registry.dimensions.conform(["instrument", "physical_filter", "detector", "exposure"])
189 self.dataset_types["raw"] = DatasetType("raw", dimensions, storage_class)
191 dimensions = registry.dimensions.conform(["instrument", "physical_filter", "detector", "visit"])
192 self.dataset_types["calexp"] = DatasetType("calexp", dimensions, storage_class)
194 dimensions = registry.dimensions.conform(["instrument", "physical_filter", "detector", "visit"])
195 self.dataset_types["no_obscore"] = DatasetType("no_obscore", dimensions, storage_class)
197 dimensions = registry.dimensions.conform(["instrument", "physical_filter", "detector"])
198 self.dataset_types["calib"] = DatasetType("calib", dimensions, storage_class, isCalibration=True)
200 for dataset_type in self.dataset_types.values():
201 registry.registerDatasetType(dataset_type)
203 # Add few run collections.
204 for run in (1, 2, 3, 4, 5, 6):
205 registry.registerRun(f"run{run}")
207 # Add few chained collections, run6 is not in any chained collections.
208 registry.registerCollection("chain12", CollectionType.CHAINED)
209 registry.setCollectionChain("chain12", ("run1", "run2"))
210 registry.registerCollection("chain34", CollectionType.CHAINED)
211 registry.setCollectionChain("chain34", ("run3", "run4"))
212 registry.registerCollection("chain-all", CollectionType.CHAINED)
213 registry.setCollectionChain("chain-all", ("chain12", "chain34", "run5"))
215 # And a tagged collection
216 registry.registerCollection("tagged", CollectionType.TAGGED)
218 def make_obscore_config(
219 self, collections: list[str] | None = None, collection_type: str | None = None
220 ) -> Config:
221 """Make configuration for obscore manager."""
222 obscore_config = Config(os.path.join(TESTDIR, "config", "basic", "obscore.yaml"))
223 if collections is not None:
224 obscore_config["collections"] = collections
225 if collection_type is not None:
226 obscore_config["collection_type"] = collection_type
227 return obscore_config
229 def _insert_dataset(
230 self, registry: SqlRegistry, run: str, dataset_type: str, do_import: bool = False, **kwargs
231 ) -> DatasetRef:
232 """Insert or import one dataset into a specified run collection."""
233 data_id = {"instrument": "DummyCam", "physical_filter": "d-r"}
234 data_id.update(kwargs)
235 coordinate = DataCoordinate.standardize(data_id, universe=registry.dimensions)
236 if do_import:
237 ds_type = self.dataset_types[dataset_type]
238 ref = DatasetRef(ds_type, coordinate, run=run)
239 [ref] = registry._importDatasets([ref])
240 else:
241 [ref] = registry.insertDatasets(dataset_type, [data_id], run=run)
242 return ref
244 def _insert_datasets(self, registry: SqlRegistry, do_import: bool = False) -> list[DatasetRef]:
245 """Inset a small bunch of datasets into every run collection."""
246 return [
247 self._insert_dataset(registry, "run1", "raw", detector=1, exposure=1, do_import=do_import),
248 self._insert_dataset(registry, "run2", "calexp", detector=2, visit=2, do_import=do_import),
249 self._insert_dataset(registry, "run3", "raw", detector=3, exposure=3, do_import=do_import),
250 self._insert_dataset(registry, "run4", "calexp", detector=4, visit=4, do_import=do_import),
251 self._insert_dataset(registry, "run5", "calexp", detector=4, visit=4, do_import=do_import),
252 # This dataset type is not configured, will not be in obscore.
253 self._insert_dataset(registry, "run5", "no_obscore", detector=1, visit=1, do_import=do_import),
254 self._insert_dataset(registry, "run6", "raw", detector=1, exposure=4, do_import=do_import),
255 ]
257 def test_config_errors(self):
258 """Test for handling various configuration problems."""
259 # This raises pydantic ValidationError, which wraps ValueError
260 exception_re = "'collections' must have one element"
261 with self.assertRaisesRegex(ValueError, exception_re):
262 self.make_registry(None, "TAGGED")
264 with self.assertRaisesRegex(ValueError, exception_re):
265 self.make_registry([], "TAGGED")
267 with self.assertRaisesRegex(ValueError, exception_re):
268 self.make_registry(["run1", "run2"], "TAGGED")
270 # Invalid regex.
271 with self.assertRaisesRegex(ValueError, "Failed to compile regex"):
272 self.make_registry(["+run"], "RUN")
274 def test_schema(self):
275 """Check how obscore schema is constructed"""
276 config = ObsCoreConfig(obs_collection="", dataset_types={}, facility_name="FACILITY")
277 schema = ObsCoreSchema(config, [])
278 table_spec = schema.table_spec
279 self.assertEqual(list(table_spec.fields.names), [col.name for col in _STATIC_COLUMNS])
281 # extra columns from top-level config
282 config = ObsCoreConfig(
283 obs_collection="",
284 extra_columns={"c1": 1, "c2": "string", "c3": {"template": "{calib_level}", "type": "float"}},
285 dataset_types={},
286 facility_name="FACILITY",
287 )
288 schema = ObsCoreSchema(config, [])
289 table_spec = schema.table_spec
290 self.assertEqual(
291 list(table_spec.fields.names),
292 [col.name for col in _STATIC_COLUMNS] + ["c1", "c2", "c3"],
293 )
294 self.assertEqual(table_spec.fields["c1"].dtype, sqlalchemy.BigInteger)
295 self.assertEqual(table_spec.fields["c2"].dtype, sqlalchemy.String)
296 self.assertEqual(table_spec.fields["c3"].dtype, sqlalchemy.Float)
298 # extra columns from per-dataset type configs
299 config = ObsCoreConfig(
300 obs_collection="",
301 extra_columns={"c1": 1},
302 dataset_types={
303 "raw": DatasetTypeConfig(
304 name="raw",
305 dataproduct_type="image",
306 calib_level=1,
307 extra_columns={"c2": "string"},
308 ),
309 "calexp": DatasetTypeConfig(
310 dataproduct_type="image",
311 calib_level=2,
312 extra_columns={"c3": 1e10},
313 ),
314 },
315 facility_name="FACILITY",
316 )
317 schema = ObsCoreSchema(config, [])
318 table_spec = schema.table_spec
319 self.assertEqual(
320 list(table_spec.fields.names),
321 [col.name for col in _STATIC_COLUMNS] + ["c1", "c2", "c3"],
322 )
323 self.assertEqual(table_spec.fields["c1"].dtype, sqlalchemy.BigInteger)
324 self.assertEqual(table_spec.fields["c2"].dtype, sqlalchemy.String)
325 self.assertEqual(table_spec.fields["c3"].dtype, sqlalchemy.Float)
327 # Columns with the same names as in static list in configs, types
328 # are not overriden.
329 config = ObsCoreConfig(
330 version=0,
331 obs_collection="",
332 extra_columns={"t_xel": 1e10},
333 dataset_types={
334 "raw": DatasetTypeConfig(
335 dataproduct_type="image",
336 calib_level=1,
337 extra_columns={"target_name": 1},
338 ),
339 "calexp": DatasetTypeConfig(
340 dataproduct_type="image",
341 calib_level=2,
342 extra_columns={"em_xel": "string"},
343 ),
344 },
345 facility_name="FACILITY",
346 )
347 schema = ObsCoreSchema(config, [])
348 table_spec = schema.table_spec
349 self.assertEqual(list(table_spec.fields.names), [col.name for col in _STATIC_COLUMNS])
350 self.assertEqual(table_spec.fields["t_xel"].dtype, sqlalchemy.Integer)
351 self.assertEqual(table_spec.fields["target_name"].dtype, sqlalchemy.String)
352 self.assertEqual(table_spec.fields["em_xel"].dtype, sqlalchemy.Integer)
354 def test_insert_existing_collection(self):
355 """Test insert and import registry methods, with various restrictions
356 on collection names.
357 """
358 # First item is collections, second item is expected record count.
359 test_data = (
360 (None, 6),
361 (["run1", "run2"], 2),
362 (["run[34]"], 2),
363 (["[rR]un[^6]"], 5),
364 )
366 for collections, count in test_data:
367 for do_import in (False, True):
368 registry = self.make_registry(collections)
369 obscore = registry.obsCoreTableManager
370 assert obscore is not None
371 self._insert_datasets(registry, do_import)
373 with obscore.query() as result:
374 rows = list(result)
375 self.assertEqual(len(rows), count)
377 # Also check `query` method with COUNT(*)
378 with obscore.query([sqlalchemy.sql.func.count()]) as result:
379 scalar = result.scalar_one()
380 self.assertEqual(scalar, count)
382 def test_drop_datasets(self):
383 """Test for dropping datasets after obscore insert."""
384 collections = None
385 registry = self.make_registry(collections)
386 obscore = registry.obsCoreTableManager
387 assert obscore is not None
388 refs = self._insert_datasets(registry)
390 with obscore.query() as result:
391 rows = list(result)
392 self.assertEqual(len(rows), 6)
394 # drop single dataset
395 registry.removeDatasets(ref for ref in refs if ref.run == "run1")
396 with obscore.query() as result:
397 rows = list(result)
398 self.assertEqual(len(rows), 5)
400 # drop whole run collection
401 registry.removeCollection("run6")
402 with obscore.query() as result:
403 rows = list(result)
404 self.assertEqual(len(rows), 4)
406 def test_associate(self):
407 """Test for associating datasets to TAGGED collection."""
408 collections = ["tagged"]
409 registry = self.make_registry(collections, "TAGGED")
410 obscore = registry.obsCoreTableManager
411 assert obscore is not None
412 refs = self._insert_datasets(registry)
414 with obscore.query() as result:
415 rows = list(result)
416 self.assertEqual(len(rows), 0)
418 # Associate datasets that are already in obscore, changes nothing.
419 registry.associate("tagged", (ref for ref in refs if ref.run == "run1"))
420 with obscore.query() as result:
421 rows = list(result)
422 self.assertEqual(len(rows), 1)
424 # Associate datasets that are not in obscore
425 registry.associate("tagged", (ref for ref in refs if ref.run == "run3"))
426 with obscore.query() as result:
427 rows = list(result)
428 self.assertEqual(len(rows), 2)
430 # Disassociate them
431 registry.disassociate("tagged", (ref for ref in refs if ref.run == "run3"))
432 with obscore.query() as result:
433 rows = list(result)
434 self.assertEqual(len(rows), 1)
436 # Non-associated dataset, should be OK and not throw.
437 registry.disassociate("tagged", (ref for ref in refs if ref.run == "run2"))
438 with obscore.query() as result:
439 rows = list(result)
440 self.assertEqual(len(rows), 1)
442 registry.disassociate("tagged", (ref for ref in refs if ref.run == "run1"))
443 with obscore.query() as result:
444 rows = list(result)
445 self.assertEqual(len(rows), 0)
447 def test_region_type_warning(self) -> None:
448 """Test that non-polygon region generates one or more warnings."""
449 collections = None
450 registry = self.make_registry(collections)
452 with self.assertWarns(RegionTypeWarning) as cm:
453 self._insert_dataset(registry, "run2", "calexp", detector=2, visit=9)
454 self.assertRegex(
455 str(cm.warning),
456 "Unexpected region type: .*lsst.sphgeom._sphgeom.Box.*",
457 )
459 def test_update_exposure_region(self) -> None:
460 """Test for update_exposure_regions method."""
461 registry = self.make_registry(["run1"])
462 obscore = registry.obsCoreTableManager
463 assert obscore is not None
465 # Exposure 4 is not associated with any visit.
466 for detector in (1, 2, 3, 4):
467 self._insert_dataset(registry, "run1", "raw", detector=detector, exposure=4)
469 # All spatial columns should be None.
470 with obscore.query() as result:
471 rows = list(result)
472 self.assertEqual(len(rows), 4)
473 for row in rows:
474 self.assertIsNone(row.s_ra)
475 self.assertIsNone(row.s_dec)
476 self.assertIsNone(row.s_region)
478 # Assign Region from visit 4.
479 count = obscore.update_exposure_regions(
480 "DummyCam", [(4, 1, self.regions[(4, 1)]), (4, 2, self.regions[(4, 2)])]
481 )
482 self.assertEqual(count, 2)
484 with obscore.query(["s_ra", "s_dec", "s_region", "lsst_detector", "facility_name"]) as result:
485 rows = list(result)
486 self.assertEqual(len(rows), 4)
487 for row in rows:
488 if row.lsst_detector in (1, 2):
489 self.assertIsNotNone(row.s_ra)
490 self.assertIsNotNone(row.s_dec)
491 self.assertIsNotNone(row.s_region)
492 else:
493 self.assertIsNone(row.s_ra)
494 self.assertIsNone(row.s_dec)
495 self.assertIsNone(row.s_region)
496 self.assertEqual(row.facility_name, "derived_facility")
499class SQLiteObsCoreTest(ObsCoreTests, unittest.TestCase):
500 """Unit test for obscore with SQLite backend."""
502 def setUp(self):
503 self.root = makeTestTempDir(TESTDIR)
505 def tearDown(self):
506 removeTestTempDir(self.root)
508 def make_registry_config(
509 self, collections: list[str] | None = None, collection_type: str | None = None
510 ) -> RegistryConfig:
511 # docstring inherited from a base class
512 _, filename = tempfile.mkstemp(dir=self.root, suffix=".sqlite3")
513 config = RegistryConfig()
514 config["db"] = f"sqlite:///{filename}"
515 config["managers", "obscore"] = {
516 "cls": "lsst.daf.butler.registry.obscore.ObsCoreLiveTableManager",
517 "config": self.make_obscore_config(collections, collection_type),
518 }
519 return config
522class ClonedSqliteObscoreTest(SQLiteObsCoreTest, unittest.TestCase):
523 """Unit test for obscore manager created via clone()"""
525 def make_registry(
526 self, collections: list[str] | None = None, collection_type: str | None = None
527 ) -> SqlRegistry:
528 """Create new empty Registry."""
529 original = super().make_registry(collections, collection_type)
530 copy = original.copy()
531 self.addCleanup(copy.close)
532 return copy
535class PostgresObsCoreTest(ObsCoreTests, unittest.TestCase):
536 """Unit test for obscore with PostgreSQL backend."""
538 @classmethod
539 def setUpClass(cls):
540 # Create the postgres test server.
541 cls.postgresql = cls.enterClassContext(setup_postgres_test_db())
542 super().setUpClass()
544 def setUp(self):
545 self.root = makeTestTempDir(TESTDIR)
547 def tearDown(self):
548 removeTestTempDir(self.root)
550 def make_registry_config(
551 self, collections: list[str] | None = None, collection_type: str | None = None
552 ) -> RegistryConfig:
553 # docstring inherited from a base class
554 config = RegistryConfig()
555 self.postgresql.patch_registry_config(config)
557 config["managers", "obscore"] = {
558 "cls": "lsst.daf.butler.registry.obscore.ObsCoreLiveTableManager",
559 "config": self.make_obscore_config(collections, collection_type),
560 }
561 return config
564class PostgresPgSphereObsCoreTest(PostgresObsCoreTest):
565 """Unit test for obscore with PostgreSQL backend and pgsphere plugin."""
567 @classmethod
568 def setUpClass(cls):
569 super().setUpClass()
570 with cls.postgresql.begin() as connection:
571 try:
572 connection.execute(sqlalchemy.text("CREATE EXTENSION pg_sphere"))
573 except sqlalchemy.exc.DatabaseError as exc:
574 raise unittest.SkipTest(f"pg_sphere extension does not exist: {exc}") from None
576 def make_obscore_config(
577 self, collections: list[str] | None = None, collection_type: str | None = None
578 ) -> Config:
579 """Make configuration for obscore manager."""
580 obscore_config = super().make_obscore_config(collections, collection_type)
581 obscore_config["spatial_plugins"] = {
582 "pgsphere": {
583 "cls": "lsst.daf.butler.registry.obscore.pgsphere.PgSphereObsCorePlugin",
584 "config": {
585 "region_column": "pgs_region",
586 "position_column": "pgs_center",
587 },
588 }
589 }
590 return obscore_config
592 def test_spatial(self):
593 """Test that pgsphere plugin fills spatial columns."""
594 collections = None
595 registry = self.make_registry(collections)
596 obscore = registry.obsCoreTableManager
597 assert obscore is not None
598 self._insert_datasets(registry)
600 # select everything
601 with obscore.query() as result:
602 rows = list(result)
603 self.assertEqual(len(rows), 6)
605 db = registry._db
606 assert registry.obsCoreTableManager is not None
607 table = cast(ObsCoreLiveTableManager, registry.obsCoreTableManager).table
609 # It's not easy to generate spatial queries in sqlalchemy, use plain
610 # text queries for testing.
612 # position matching visit=1, there is a single dataset
613 query = f"SELECT * FROM {table.key} WHERE pgs_center <-> '(2d,0d)'::spoint < .1"
614 with db.query(sqlalchemy.text(query)) as results:
615 self.assertEqual(len(list(results)), 1)
617 # position matching visit=4, there are two datasets
618 query = f"SELECT * FROM {table.key} WHERE pgs_center <-> '(272d,0d)'::spoint < .1"
619 with db.query(sqlalchemy.text(query)) as results:
620 self.assertEqual(len(list(results)), 2)
622 # position matching visit=1, there is a single dataset
623 query = f"SELECT * FROM {table.key} WHERE '(2d,-3d)'::spoint @ pgs_region"
624 with db.query(sqlalchemy.text(query)) as results:
625 self.assertEqual(len(list(results)), 1)
627 # position matching visit=4, there are two datasets
628 query = f"SELECT * FROM {table.key} WHERE '(272d,3d)'::spoint @ pgs_region"
629 with db.query(sqlalchemy.text(query)) as results:
630 self.assertEqual(len(list(results)), 2)
633if __name__ == "__main__":
634 unittest.main()