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