Skip to content

Commit

Permalink
Add tests for storageclass conversion on get
Browse files Browse the repository at this point in the history
  • Loading branch information
timj committed Dec 10, 2021
1 parent 59723ce commit c68f4b8
Show file tree
Hide file tree
Showing 3 changed files with 59 additions and 1 deletion.
19 changes: 18 additions & 1 deletion python/lsst/daf/butler/tests/_examplePythonTypes.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,14 @@
formats such as FITS or HDF5.
"""

__all__ = ("ListDelegate", "MetricsDelegate", "MetricsExample", "registerMetricsExample")
__all__ = ("ListDelegate", "MetricsDelegate", "MetricsExample", "registerMetricsExample",
"MetricsExampleModel")


import copy
from typing import Optional, Any, Dict, List

from pydantic import BaseModel
from lsst.daf.butler import StorageClassDelegate, StorageClass


Expand Down Expand Up @@ -234,6 +238,19 @@ def makeFromDict(cls, exportDict):
return cls(exportDict["summary"], exportDict["output"], data)


class MetricsExampleModel(BaseModel):
"""A variant of `MetricsExample` based on model."""

summary: Optional[Dict[str, Any]]
output: Optional[Dict[str, Any]]
data: Optional[List[Any]]

@classmethod
def from_metrics(cls, metrics: MetricsExample) -> "MetricsExampleModel":
"""Create a model based on an example."""
return cls.parse_obj(metrics.exportAsDict())


class ListDelegate(StorageClassDelegate):
"""Parameter handler for list parameters"""

Expand Down
5 changes: 5 additions & 0 deletions tests/config/basic/storageClasses.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -79,3 +79,8 @@ storageClasses:
counter: Integer
StructuredCompositeReadCompNoDisassembly:
inheritsFrom: StructuredCompositeReadComp
StructuredDataNoComponentsModel:
# Reading and writing a blob and no components known
pytype: lsst.daf.butler.tests.MetricsExampleModel
converters:
lsst.daf.butler.tests.MetricsExample: lsst.daf.butler.tests.MetricsExampleModel.from_metrics
36 changes: 36 additions & 0 deletions tests/test_butler.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,7 @@ def mock_s3(cls):
from threading import Thread
from tempfile import gettempdir
from lsst.utils import doImport
from lsst.utils.introspection import get_full_type_name
from lsst.daf.butler import Butler, Config, ButlerConfig
from lsst.daf.butler import StorageClassFactory
from lsst.daf.butler import DatasetType, DatasetRef, DatasetIdGenEnum
Expand Down Expand Up @@ -1286,6 +1287,41 @@ def testPruneDatasets(self):
# Clear out the datasets from registry.
butler.pruneDatasets([ref1, ref2, ref3], purge=True, unstore=True)

def testPytypeCoercion(self):
"""Test python type coercion on Butler.get"""

# Store some data with the normal example storage class.
storageClass = self.storageClassFactory.getStorageClass("StructuredDataNoComponents")
datasetTypeName = "test_metric"
butler = self.runPutGetTest(storageClass, datasetTypeName)

dataId = {"instrument": "DummyCamComp", "visit": 423}
metric = butler.get(datasetTypeName, dataId=dataId)
self.assertEqual(get_full_type_name(metric), "lsst.daf.butler.tests.MetricsExample")

datasetType_ori = butler.registry.getDatasetType(datasetTypeName)
self.assertEqual(datasetType_ori.storageClass.name, "StructuredDataNoComponents")

# Now need to hack the registry dataset type definition.
# There is no API for this.
manager = butler.registry._managers.datasets
manager._db.update(
manager._static.dataset_type,
{"name": datasetTypeName},
{datasetTypeName: datasetTypeName, "storage_class": "StructuredDataNoComponentsModel"},
)

# Force reset od dataset type cache
butler.registry.refresh()

datasetType_new = butler.registry.getDatasetType(datasetTypeName)
self.assertEqual(datasetType_new.name, datasetType_ori.name)
self.assertEqual(datasetType_new.storageClass.name, "StructuredDataNoComponentsModel")

metric_model = butler.get(datasetTypeName, dataId=dataId)
self.assertNotEqual(type(metric_model), type(metric))
self.assertEqual(get_full_type_name(metric_model), "lsst.daf.butler.tests.MetricsExampleModel")


class InMemoryDatastoreButlerTestCase(ButlerTests, unittest.TestCase):
"""InMemoryDatastore specialization of a butler"""
Expand Down

0 comments on commit c68f4b8

Please sign in to comment.