Add test for FormatterV2.add_provenance

python/lsst/daf/butler/tests/_examplePythonTypes.py

@@ -46,10 +46,11 @@
 import copy
 import dataclasses
 import types
+import uuid
 from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any
 
-from lsst.daf.butler import StorageClass, StorageClassDelegate
+from lsst.daf.butler import DatasetProvenance, StorageClass, StorageClassDelegate
 from pydantic import BaseModel
 
 if TYPE_CHECKING:
@@ -297,6 +298,8 @@ class MetricsExampleModel(BaseModel):
     summary: dict[str, Any] | None = None
     output: dict[str, Any] | None = None
     data: list[Any] | None = None
+    provenance: DatasetProvenance | None = None
+    dataset_id: uuid.UUID | None = None
 
     @classmethod
     def from_metrics(cls, metrics: MetricsExample) -> MetricsExampleModel:

python/lsst/daf/butler/tests/testFormatters.py

@@ -44,6 +44,7 @@
 from lsst.resources import ResourceHandleProtocol
 
 from .._formatter import Formatter, FormatterV2
+from ..formatters.json import JsonFormatter
 from ..formatters.yaml import YamlFormatter
 
 if TYPE_CHECKING:
@@ -271,3 +272,14 @@ def write(self, inMemoryDataset: Any) -> None:
 
         with open(fileDescriptor.location.path, "w") as fd:
             yaml.dump(inMemoryDataset, fd)
+
+
+class MetricsExampleModelProvenanceFormatter(JsonFormatter):
+    """Specialist formatter to test provenance addition."""
+
+    def add_provenance(self, in_memory_dataset: Any) -> Any:
+        # Copy it to prove that works.
+        new = in_memory_dataset.model_copy()
+        new.provenance = self._provenance
+        new.dataset_id = self.dataset_ref.id
+        return new

tests/config/basic/formatters.yaml

@@ -47,3 +47,4 @@ MetricsExampleModelB: lsst.daf.butler.formatters.yaml.YamlFormatter
 TupleExampleA: lsst.daf.butler.formatters.json.JsonFormatter
 TupleExampleB: lsst.daf.butler.formatters.yaml.YamlFormatter
 MetricsConversion: lsst.daf.butler.formatters.json.JsonFormatter
+MetricsExampleModelProvenance: lsst.daf.butler.tests.testFormatters.MetricsExampleModelProvenanceFormatter

tests/config/basic/storageClasses.yaml

@@ -145,3 +145,5 @@ storageClasses:
     pytype: tuple
   TupleExampleB:
     pytype: tuple
+  MetricsExampleModelProvenance:
+    pytype: lsst.daf.butler.tests.MetricsExampleModel

tests/test_butler.py

@@ -77,6 +77,7 @@ def mock_aws(*args: Any, **kwargs: Any) -> Any:  # type: ignore[no-untyped-def]
     DataCoordinate,
     DatasetExistence,
     DatasetNotFoundError,
+    DatasetProvenance,
     DatasetRef,
     DatasetType,
     FileDataset,
@@ -101,7 +102,7 @@ def mock_aws(*args: Any, **kwargs: Any) -> Any:  # type: ignore[no-untyped-def]
 )
 from lsst.daf.butler.registry.sql_registry import SqlRegistry
 from lsst.daf.butler.repo_relocation import BUTLER_ROOT_TAG
-from lsst.daf.butler.tests import MetricsExample, MultiDetectorFormatter
+from lsst.daf.butler.tests import MetricsExample, MetricsExampleModel, MultiDetectorFormatter
 from lsst.daf.butler.tests.postgresql import TemporaryPostgresInstance, setup_postgres_test_db
 from lsst.daf.butler.tests.utils import TestCaseMixin, makeTestTempDir, removeTestTempDir, safeTestTempDir
 from lsst.resources import ResourcePath
@@ -2124,6 +2125,33 @@ def testPytypeCoercion(self) -> None:
         with self.assertRaises(ValueError):
             butler.get(datasetTypeName, dataId=dataId)
 
+    def test_provenance(self):
+        """Test that provenance is attached on put."""
+        run = "test_run"
+        butler, datasetType = self.create_butler(run, "MetricsExampleModelProvenance", "prov_metric")
+        metric = MetricsExampleModel(
+            summary={"AM1": 5.2, "AM2": 30.6},
+            output={"a": [1, 2, 3], "b": {"blue": 5, "red": "green"}},
+            data=[563, 234, 456.7, 752, 8, 9, 27],
+        )
+        # Provenance can be attached to the object being put. Whether
+        # it is or not is dependent on the formatter. For this test we
+        # copy on adding provenance to ensure they differ.
+        self.assertIsNone(metric.dataset_id)
+        metric_ref = butler.put(metric, datasetType, visit=424, instrument="DummyCamComp")
+        self.assertIsNone(metric.dataset_id)
+        metric_2 = butler.get(metric_ref)
+        self.assertEqual(metric_2.data, metric.data)
+        self.assertEqual(metric_2.dataset_id, metric_ref.id)
+        self.assertIsNone(metric_2.provenance)
+
+        # Put with provenance.
+        prov = DatasetProvenance(quantum_id=uuid.uuid4())
+        prov.add_input(metric_ref)
+        metric_ref2 = butler.put(metric, datasetType, visit=423, instrument="DummyCamComp", provenance=prov)
+        metric_3 = butler.get(metric_ref2)
+        self.assertEqual(metric_3.provenance, prov)
+
 
 class PostgresPosixDatastoreButlerTestCase(FileDatastoreButlerTests, unittest.TestCase):
     """PosixDatastore specialization of a butler using Postgres"""