Improves Rules model - use enums where I know what the values mean, Literals where I don't. Also deserialise dates properly.

Author: acerathereinspirative-nhs-nvir

src/eligibility_signposting_api/model/rules.py

@@ -1,11 +1,48 @@
 from __future__ import annotations
 
-from typing import NewType
+import typing
+from datetime import date, datetime
+from enum import Enum
+from typing import Literal, NewType
 
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, Field, field_serializer, field_validator
+
+if typing.TYPE_CHECKING:
+    from pydantic import SerializationInfo
 
-Campaign = NewType("Campaign", str)
 BucketName = NewType("BucketName", str)
+CampaignName = NewType("CampaignName", str)
+CampaignVersion = NewType("CampaignVersion", str)
+CampaignID = NewType("CampaignID", str)
+IterationName = NewType("IterationName", str)
+IterationVersion = NewType("IterationVersion", str)
+IterationID = NewType("IterationID", str)
+RuleName = NewType("RuleName", str)
+RuleDescription = NewType("RuleDescription", str)
+RulePriority = NewType("RulePriority", int)
+RuleAttributeLevel = NewType("RuleAttributeLevel", str)
+RuleAttributeName = NewType("RuleAttributeName", str)
+RuleComparator = NewType("RuleComparator", str)
+StartDate = NewType("StartDate", date)
+EndDate = NewType("EndDate", date)
+
+
+class RuleType(str, Enum):
+    filter = "F"
+    suppression = "S"
+    redirect = "R"
+
+
+class RuleOperator(str, Enum):
+    lt = "<"
+    gt = ">"
+    year_gt = "Y>"
+    not_in = "not_in"
+    equals = "="
+    lte = "<="
+    ne = "!="
+    date_gte = "D>="
+    member_of = "MemberOf"
 
 
 class IterationCohort(BaseModel):
@@ -16,57 +53,72 @@ class IterationCohort(BaseModel):
 
 
 class IterationRule(BaseModel):
-    type: str | None = Field(None, alias="Type")
-    name: str | None = Field(None, alias="Name")
-    description: str | None = Field(None, alias="Description")
-    priority: int | None = Field(None, alias="Priority")
-    attribute_level: str | None = Field(None, alias="AttributeLevel")
-    attribute_name: str | None = Field(None, alias="AttributeName")
-    operator: str | None = Field(None, alias="Operator")
-    comparator: str | None = Field(None, alias="Comparator")
+    type: RuleType = Field(..., alias="Type")
+    name: RuleName = Field(..., alias="Name")
+    description: RuleDescription = Field(..., alias="Description")
+    priority: RulePriority = Field(..., alias="Priority")
+    attribute_level: RuleAttributeLevel = Field(..., alias="AttributeLevel")
+    attribute_name: RuleAttributeName = Field(..., alias="AttributeName")
+    operator: RuleOperator = Field(..., alias="Operator")
+    comparator: RuleComparator = Field(..., alias="Comparator")
     attribute_target: str | None = Field(None, alias="AttributeTarget")
     comms_routing: str | None = Field(None, alias="CommsRouting")
 
     model_config = {"populate_by_name": True}
 
 
 class Iteration(BaseModel):
-    id: str = Field(..., alias="ID")
+    id: IterationID = Field(..., alias="ID")
     default_comms_routing: str | None = Field(None, alias="DefaultCommsRouting")
-    version: int | None = Field(None, alias="Version")
-    name: str | None = Field(None, alias="Name")
+    version: IterationVersion = Field(..., alias="Version")
+    name: IterationName = Field(..., alias="Name")
     iteration_date: str | None = Field(None, alias="IterationDate")
     iteration_number: int | None = Field(None, alias="IterationNumber")
-    comms_type: str | None = Field(None, alias="CommsType")
+    comms_type: Literal["I", "R"] = Field(..., alias="CommsType")
     approval_minimum: int | None = Field(None, alias="ApprovalMinimum")
     approval_maximum: int | None = Field(None, alias="ApprovalMaximum")
-    type: str | None = Field(None, alias="Type")
+    type: Literal["A", "M", "S"] = Field(..., alias="Type")
     iteration_cohorts: list[IterationCohort] | None = Field(None, alias="IterationCohorts")
     iteration_rules: list[IterationRule] | None = Field(None, alias="IterationRules")
 
     model_config = {"populate_by_name": True}
 
 
 class CampaignConfig(BaseModel):
-    id: str = Field(..., alias="ID")
-    version: int = Field(..., alias="Version")
-    name: str = Field(..., alias="Name")
-    type: str | None = Field(None, alias="Type")
-    target: str | None = Field(None, alias="Target")
+    id: CampaignID = Field(..., alias="ID")
+    version: CampaignVersion = Field(..., alias="Version")
+    name: CampaignName = Field(..., alias="Name")
+    type: Literal["V", "S"] = Field(..., alias="Type")
+    target: Literal["COVID", "FLU", "MMR", "RSV"] = Field(..., alias="Target")
     manager: str | None = Field(None, alias="Manager")
     approver: str | None = Field(None, alias="Approver")
     reviewer: str | None = Field(None, alias="Reviewer")
-    iteration_frequency: str | None = Field(None, alias="IterationFrequency")
-    iteration_type: str | None = Field(None, alias="IterationType")
+    iteration_frequency: Literal["X", "D", "W", "M", "Q", "A"] = Field(..., alias="IterationFrequency")
+    iteration_type: Literal["A", "M", "S"] = Field(..., alias="IterationType")
     iteration_time: str | None = Field(None, alias="IterationTime")
     default_comms_routing: str | None = Field(None, alias="DefaultCommsRouting")
-    start_date: str | None = Field(None, alias="StartDate")
-    end_date: str | None = Field(None, alias="EndDate")
+    start_date: StartDate = Field(..., alias="StartDate")
+    end_date: EndDate = Field(..., alias="EndDate")
     approval_minimum: int | None = Field(None, alias="ApprovalMinimum")
     approval_maximum: int | None = Field(None, alias="ApprovalMaximum")
     iterations: list[Iteration] | None = Field(None, alias="Iterations")
 
-    model_config = {"populate_by_name": True}
+    model_config = {
+        "populate_by_name": True,
+        "arbitrary_types_allowed": True,
+    }
+
+    @field_validator("start_date", "end_date", mode="before")
+    @classmethod
+    def parse_dates(cls, v: str | date) -> date:
+        if isinstance(v, date):
+            return v
+        return datetime.strptime(v, "%Y%m%d").date()  # noqa: DTZ007
+
+    @field_serializer("start_date", "end_date", when_used="always")
+    @staticmethod
+    def serialize_dates(v: date, _info: SerializationInfo) -> str:
+        return v.strftime("%Y%m%d")
 
 
 class Rules(BaseModel):

src/eligibility_signposting_api/repos/rules_repo.py

@@ -4,7 +4,7 @@
 from botocore.client import BaseClient
 from wireup import Inject, service
 
-from eligibility_signposting_api.model.rules import BucketName, Campaign, CampaignConfig, Rules
+from eligibility_signposting_api.model.rules import BucketName, CampaignConfig, CampaignName, Rules
 
 
 @service
@@ -18,7 +18,7 @@ def __init__(
         self.s3_client = s3_client
         self.bucket_name = bucket_name
 
-    def get_campaign_config(self, campaign: Campaign) -> CampaignConfig:
+    def get_campaign_config(self, campaign: CampaignName) -> CampaignConfig:
         response = self.s3_client.get_object(Bucket=self.bucket_name, Key=f"{campaign}.json")
         body = response["Body"].read()
         return Rules.model_validate(json.loads(body)).campaign_config

tests/integration/repo/test_rules_repo.py

@@ -9,7 +9,7 @@
 from eligibility_signposting_api.repos.rules_repo import RulesRepo
 from tests.integration.conftest import AWS_REGION
 from tests.utils.builders import CampaignConfigFactory, random_str
-from tests.utils.rules.campaign import is_campaign_config
+from tests.utils.rules.rules import is_campaign_config
 
 
 @pytest.fixture

tests/utils/builders.py

@@ -17,8 +17,8 @@ class IterationRuleFactory(ModelFactory[IterationRule]):
 
 class IterationFactory(ModelFactory[Iteration]):
     __model__ = Iteration
-    iteration_cohorts: Use(IterationCohortFactory.batch, size=2)
-    iteration_rules: Use(IterationRuleFactory.batch, size=2)
+    iteration_cohorts = Use(IterationCohortFactory.batch, size=2)
+    iteration_rules = Use(IterationRuleFactory.batch, size=2)
 
 
 class CampaignConfigFactory(ModelFactory[CampaignConfig]):

tests/utils/rules/campaign.py tests/utils/rules/rules.pytests/utils/rules/campaign.py renamed to tests/utils/rules/rules.py

@@ -5,17 +5,17 @@
 from hamcrest.core.helpers.wrap_matcher import wrap_matcher
 from hamcrest.core.matcher import Matcher
 
-from eligibility_signposting_api.model.rules import CampaignConfig
+from eligibility_signposting_api.model.rules import CampaignConfig, CampaignID, CampaignName, CampaignVersion
 
 ANYTHING = anything()
 
 
 class CampaignConfigMatcher(BaseMatcher[CampaignConfig]):
     def __init__(self):
         super().__init__()
-        self.id_: Matcher[str] = ANYTHING
-        self.name: Matcher[str] = ANYTHING
-        self.version: Matcher[str] = ANYTHING
+        self.id_: Matcher[CampaignID] = ANYTHING
+        self.name: Matcher[CampaignName] = ANYTHING
+        self.version: Matcher[CampaignVersion] = ANYTHING
 
     def describe_to(self, description: Description) -> None:
         description.append_text("CampaignConfig with")
@@ -38,25 +38,25 @@ def describe_match(self, item: CampaignConfig, match_description: Description) -
         describe_field_match(self.name, "name", item.name, match_description)
         describe_field_match(self.version, "version", item.version, match_description)
 
-    def with_id(self, id_: str | Matcher[str]):
+    def with_id(self, id_: CampaignID | Matcher[CampaignID]):
         self.id_ = wrap_matcher(id_)
         return self
 
-    def and_id(self, id_: str | Matcher[str]):
+    def and_id(self, id_: CampaignID | Matcher[CampaignID]):
         return self.with_id(id_)
 
-    def with_name(self, name: str | Matcher[str]):
+    def with_name(self, name: CampaignName | Matcher[CampaignName]):
         self.name = wrap_matcher(name)
         return self
 
-    def and_name(self, name: str | Matcher[str]):
+    def and_name(self, name: CampaignName | Matcher[CampaignName]):
         return self.with_name(name)
 
-    def with_version(self, version: str | Matcher[str]):
+    def with_version(self, version: CampaignVersion | Matcher[CampaignVersion]):
         self.version = wrap_matcher(version)
         return self
 
-    def and_version(self, version: str | Matcher[str]):
+    def and_version(self, version: CampaignVersion | Matcher[CampaignVersion]):
         return self.with_version(version)