bugfix(serde): Refactor validation function for enum fields

Author: tcjennings

src/lsst/cmservice/db/manifests_v2.py

@@ -1,23 +1,19 @@
-from typing import Annotated
+"""Module for models representing generic CM Service manifests."""
 
 from pydantic import AliasChoices
 from sqlmodel import Field, SQLModel
 
 from ..common.enums import ManifestKind
-from ..models.serde import EnumSerializer, ManifestKindEnumValidator
+from ..common.types import KindField
 
 
-# this can probably be a BaseModel since this is not a db relation, but the
-# distinction probably doesn't matter
 class ManifestWrapper(SQLModel):
-    """a model for an object's Manifest wrapper, used by APIs where the `spec`
+    """A model for an object's Manifest wrapper, used by APIs where the `spec`
     should be the kind's table model, more or less.
     """
 
     apiversion: str = Field(default="io.lsst.cmservice/v1")
-    kind: Annotated[ManifestKind, ManifestKindEnumValidator, EnumSerializer] = Field(
-        default=ManifestKind.other,
-    )
+    kind: KindField = Field(default=ManifestKind.other)
     metadata_: dict = Field(
         default_factory=dict,
         schema_extra={

src/lsst/cmservice/models/serde.py

@@ -2,10 +2,27 @@
 other derivative models.
 """
 
+from enum import EnumType
+from functools import partial
+
 from pydantic import PlainSerializer, PlainValidator
 
 from ..common.enums import ManifestKind, StatusEnum
 
+
+def EnumValidator[T: EnumType](value: str | int, enum_: T) -> T:
+    """Create an enum from the input value. The input can be either the
+    enum name or its value.
+
+    Used as a Validator for a pydantic field.
+    """
+    try:
+        new_enum: T = enum_[value] if value in enum_.__members__ else enum_(value)
+    except (KeyError, ValueError):
+        raise ValueError(f"Value must be a member of {enum_.__qualname__}")
+    return new_enum
+
+
 EnumSerializer = PlainSerializer(
     lambda x: x.name,
     return_type="str",
@@ -14,15 +31,13 @@
 """A serializer for enums that produces its name, not the value."""
 
 
-StatusEnumValidator = PlainValidator(lambda x: StatusEnum[x] if isinstance(x, str) else StatusEnum(x))
+StatusEnumValidator = PlainValidator(partial(EnumValidator, enum_=StatusEnum))
 """A validator for the StatusEnum that can parse the enum from either a name
 or a value.
 """
 
 
-ManifestKindEnumValidator = PlainValidator(
-    lambda x: ManifestKind[x] if isinstance(x, str) else ManifestKind(x)
-)
+ManifestKindEnumValidator = PlainValidator(partial(EnumValidator, enum_=ManifestKind))
 """A validator for the ManifestKindEnum that can parse the enum from a name
 or a value.
 """

tests/models/test_serde.py

@@ -0,0 +1,33 @@
+import pytest
+from pydantic import BaseModel, ValidationError
+
+from lsst.cmservice.common.enums import ManifestKind, StatusEnum
+from lsst.cmservice.common.types import KindField, StatusField
+
+
+class TestModel(BaseModel):
+    status: StatusField
+    kind: KindField
+
+
+def test_validators() -> None:
+    """Test model field enum validators."""
+    # test enum validation by name and value
+    x = TestModel(status=0, kind="campaign")
+    assert x.status is StatusEnum.waiting
+    assert x.kind is ManifestKind.campaign
+
+    # test bad input (wrong name)
+    with pytest.raises(ValidationError):
+        x = TestModel(status="bad", kind="edge")
+
+    # test bad input (bad value)
+    with pytest.raises(ValidationError):
+        x = TestModel(status="waiting", kind=99)
+
+
+def test_serializers() -> None:
+    x = TestModel(status="accepted", kind="node")
+    y = x.model_dump()
+    assert y["status"] == "accepted"
+    assert y["kind"] == "node"