VED-898: Refactor RecordAttributes so it handles Immunization FHIR object rather than a dictionary (#1259)

* Refactor RecordAttributes to use Immunization object and update related methods in FhirService and tests. This change enhances type safety and improves the handling of immunization data throughout the repository and service layers.

* chore: empty commit

EmptyCommit:

* Refactor FhirService and validation_utils: Rename variable for clarity and enhance get_vaccine_type to accept Immunization model. Update references accordingly.

Author: Thomas-Boyle

lambdas/backend/src/repository/fhir_repository.py

@@ -66,27 +66,30 @@ class RecordAttributes:
     pk: str
     patient_pk: str
     patient_sk: str
-    resource: dict
     patient: dict
     vaccine_type: str
     timestamp: int
     identifier: str
-
-    def __init__(self, imms: dict, patient: any):
-        """Create attributes that may be used in dynamodb table"""
-        imms_id = imms["id"]
-        self.pk = _make_immunization_pk(imms_id)
-        if patient or imms:
-            nhs_number = get_nhs_number(imms)
-        self.patient_pk = _make_patient_pk(nhs_number)
-        self.patient = patient
-        self.resource = imms
-        self.timestamp = int(time.time())
-        self.vaccine_type = get_vaccine_type(imms)
-        self.system_id = imms["identifier"][0]["system"]
-        self.system_value = imms["identifier"][0]["value"]
-        self.patient_sk = f"{self.vaccine_type}#{imms_id}"
-        self.identifier = f"{self.system_id}#{self.system_value}"
+    immunization: Immunization
+
+    @classmethod
+    def from_immunization(cls, immunization: Immunization, patient: dict | None = None) -> "RecordAttributes":
+        """Build DynamoDB attributes from a FHIR Immunization resource."""
+        imms_dict = immunization.dict()
+        patient_resolved = patient if patient is not None else get_contained_patient(imms_dict)
+        nhs_number = get_nhs_number(imms_dict)
+        vaccine_type = get_vaccine_type(imms_dict)
+        first_identifier = immunization.identifier[0]
+        return cls(
+            pk=_make_immunization_pk(immunization.id),
+            patient_pk=_make_patient_pk(nhs_number),
+            patient_sk=f"{vaccine_type}#{immunization.id}",
+            patient=patient_resolved,
+            vaccine_type=vaccine_type,
+            timestamp=int(time.time()),
+            identifier=f"{first_identifier.system}#{first_identifier.value}",
+            immunization=immunization,
+        )
 
 
 class ImmunizationRepository:
@@ -162,10 +165,7 @@ def check_immunization_identifier_exists(self, system: str, unique_id: str) -> b
 
     def create_immunization(self, immunization: Immunization, supplier_system: str) -> Id:
         """Creates a new immunization record returning the unique id if successful."""
-        immunization_as_dict = immunization.dict()
-
-        patient = get_contained_patient(immunization_as_dict)
-        attr = RecordAttributes(immunization_as_dict, patient)
+        attr = RecordAttributes.from_immunization(immunization)
 
         response = self.table.put_item(
             Item={
@@ -188,13 +188,11 @@ def create_immunization(self, immunization: Immunization, supplier_system: str)
     def update_immunization(
         self,
         imms_id: str,
-        immunization: dict,
+        immunization: Immunization,
         existing_record_meta: ImmunizationRecordMetadata,
         supplier_system: str,
     ) -> int:
-        # VED-898 - consider refactoring to pass FHIR Immunization object rather than dict between Service -> Repository
-        patient = get_contained_patient(immunization)
-        attr = RecordAttributes(immunization, patient)
+        attr = RecordAttributes.from_immunization(immunization)
         reinstate_operation_required = existing_record_meta.is_deleted
 
         update_exp = self._build_update_expression(is_reinstate=reinstate_operation_required)
@@ -245,7 +243,7 @@ def _perform_dynamo_update(
             ":timestamp": attr.timestamp,
             ":patient_pk": attr.patient_pk,
             ":patient_sk": attr.patient_sk,
-            ":imms_resource_val": json.dumps(attr.resource, use_decimal=True),
+            ":imms_resource_val": attr.immunization.json(use_decimal=True),
             ":operation": "UPDATE",
             ":version": updated_version,
             ":supplier_system": supplier_system,

lambdas/backend/src/service/fhir_service.py

@@ -143,34 +143,31 @@ def update_immunization(self, imms_id: str, immunization: dict, supplier_system:
         except (ValueError, MandatoryError) as error:
             raise CustomValidationError(message=str(error)) from error
 
+        immunization_to_update = Immunization.parse_obj(immunization)
+
         existing_immunization_resource, existing_immunization_meta = (
             self.immunization_repo.get_immunization_resource_and_metadata_by_id(imms_id, include_deleted=True)
         )
 
         if not existing_immunization_resource:
             raise ResourceNotFoundError(resource_type="Immunization", resource_id=imms_id)
 
-        # If the user is updating the resource vaccination_type, they must have permissions for both the existing and
-        # new type. In most cases it will be the same, but it is possible for users to update the vacc type
+        existing_immunization = Immunization.parse_obj(existing_immunization_resource)
+
         if not self.authoriser.authorise(
             supplier_system,
             ApiOperationCode.UPDATE,
-            {get_vaccine_type(immunization), get_vaccine_type(existing_immunization_resource)},
+            {get_vaccine_type(immunization_to_update), get_vaccine_type(existing_immunization)},
         ):
             raise UnauthorizedVaxError()
 
-        identifier = Identifier.construct(
-            system=immunization["identifier"][0]["system"],
-            value=immunization["identifier"][0]["value"],
-        )
-
-        validate_identifiers_match(identifier, existing_immunization_meta.identifier)
+        validate_identifiers_match(immunization_to_update.identifier[0], existing_immunization_meta.identifier)
 
         if not existing_immunization_meta.is_deleted:
             validate_resource_versions_match(resource_version, existing_immunization_meta.resource_version, imms_id)
 
         return self.immunization_repo.update_immunization(
-            imms_id, immunization, existing_immunization_meta, supplier_system
+            imms_id, immunization_to_update, existing_immunization_meta, supplier_system
         )
 
     def delete_immunization(self, imms_id: str, supplier_system: str) -> None:

lambdas/backend/tests/repository/test_fhir_repository.py

@@ -391,6 +391,7 @@ def test_update_immunisation_is_successful(self):
         """it should update the immunisation record"""
         imms_id = "an-imms-id"
         imms = create_covid_immunization_dict(imms_id, VALID_NHS_NUMBER)
+        immunization = Immunization.parse_obj(imms)
         identifier = Identifier(system=imms["identifier"][0]["system"], value=imms["identifier"][0]["value"])
         existing_record_metadata = ImmunizationRecordMetadata(
             identifier=identifier, resource_version=1, is_deleted=False, is_reinstated=False
@@ -400,7 +401,7 @@ def test_update_immunisation_is_successful(self):
         self.table.update_item = MagicMock(return_value=dynamo_response)
 
         # When
-        updated_version = self.repository.update_immunization(imms_id, imms, existing_record_metadata, "Test")
+        updated_version = self.repository.update_immunization(imms_id, immunization, existing_record_metadata, "Test")
 
         # Then
         update_exp = (
@@ -420,7 +421,7 @@ def test_update_immunisation_is_successful(self):
                 ":timestamp": ANY,
                 ":patient_pk": _make_patient_pk(patient_id),
                 ":patient_sk": patient_sk,
-                ":imms_resource_val": json.dumps(imms),
+                ":imms_resource_val": immunization.json(use_decimal=True),
                 ":operation": "UPDATE",
                 ":version": 2,
                 ":supplier_system": "Test",
@@ -433,6 +434,7 @@ def test_update_immunisation_is_successful_when_record_needs_to_be_reinstated(se
         """it should reinstate a deleted record when requested via the update operation"""
         imms_id = "an-imms-id"
         imms = create_covid_immunization_dict(imms_id, VALID_NHS_NUMBER)
+        immunization = Immunization.parse_obj(imms)
         identifier = Identifier(system=imms["identifier"][0]["system"], value=imms["identifier"][0]["value"])
         existing_record_metadata = ImmunizationRecordMetadata(
             identifier=identifier, resource_version=2, is_deleted=True, is_reinstated=False
@@ -442,7 +444,7 @@ def test_update_immunisation_is_successful_when_record_needs_to_be_reinstated(se
         self.table.update_item = MagicMock(return_value=dynamo_response)
 
         # When
-        updated_version = self.repository.update_immunization(imms_id, imms, existing_record_metadata, "Test")
+        updated_version = self.repository.update_immunization(imms_id, immunization, existing_record_metadata, "Test")
 
         # Then
         update_exp = (
@@ -462,7 +464,7 @@ def test_update_immunisation_is_successful_when_record_needs_to_be_reinstated(se
                 ":timestamp": ANY,
                 ":patient_pk": _make_patient_pk(patient_id),
                 ":patient_sk": patient_sk,
-                ":imms_resource_val": json.dumps(imms),
+                ":imms_resource_val": immunization.json(use_decimal=True),
                 ":operation": "UPDATE",
                 ":version": 3,
                 ":supplier_system": "Test",
@@ -477,6 +479,7 @@ def test_update_throws_error_when_response_can_not_be_handled(self):
         condition, as a check is made first to retrieve the record."""
         imms_id = "an-id"
         imms = create_covid_immunization_dict(imms_id, VALID_NHS_NUMBER)
+        immunization = Immunization.parse_obj(imms)
         identifier = Identifier(system=imms["identifier"][0]["system"], value=imms["identifier"][0]["value"])
         existing_record_metadata = ImmunizationRecordMetadata(
             identifier=identifier, resource_version=2, is_deleted=True, is_reinstated=False
@@ -489,7 +492,7 @@ def test_update_throws_error_when_response_can_not_be_handled(self):
 
         with self.assertRaises(ResourceNotFoundError) as e:
             # When
-            self.repository.update_immunization(imms_id, imms, existing_record_metadata, "Test")
+            self.repository.update_immunization(imms_id, immunization, existing_record_metadata, "Test")
 
         # Then
         self.assertEqual(str(e.exception), "Immunization resource does not exist. ID: an-id")
@@ -728,9 +731,10 @@ def run_update_immunization_test(self, imms_id, imms, updated_dose_quantity=None
         existing_record_metadata = ImmunizationRecordMetadata(
             identifier=identifier, resource_version=1, is_deleted=False, is_reinstated=False
         )
+        immunization = Immunization.parse_obj(imms)
 
         # When
-        updated_version = self.repository.update_immunization(imms_id, imms, existing_record_metadata, "Test")
+        updated_version = self.repository.update_immunization(imms_id, immunization, existing_record_metadata, "Test")
         self.assertEqual(updated_version, 2)
 
         update_exp = (
@@ -750,7 +754,7 @@ def run_update_immunization_test(self, imms_id, imms, updated_dose_quantity=None
                 ":timestamp": ANY,
                 ":patient_pk": _make_patient_pk(patient_id),
                 ":patient_sk": patient_sk,
-                ":imms_resource_val": json.dumps(imms, use_decimal=True),
+                ":imms_resource_val": immunization.json(use_decimal=True),
                 ":operation": "UPDATE",
                 ":version": 2,
                 ":supplier_system": "Test",

lambdas/backend/tests/service/test_fhir_service.py

@@ -529,9 +529,13 @@ def test_update_immunization(self):
         self.imms_repo.get_immunization_resource_and_metadata_by_id.assert_called_once_with(
             imms_id, include_deleted=True
         )
-        self.imms_repo.update_immunization.assert_called_once_with(
-            imms_id, updated_immunisation, existing_resource_meta, "Test"
-        )
+        self.imms_repo.update_immunization.assert_called_once()
+        call_args = self.imms_repo.update_immunization.call_args[0]
+        self.assertEqual(call_args[0], imms_id)
+        self.assertIsInstance(call_args[1], Immunization)
+        self.assertEqual(call_args[1].id, imms_id)
+        self.assertEqual(call_args[2], existing_resource_meta)
+        self.assertEqual(call_args[3], "Test")
         self.authoriser.authorise.assert_called_once_with("Test", ApiOperationCode.UPDATE, {"COVID"})
 
     def test_update_immunization_raises_validation_exception_when_nhs_number_invalid(self):

lambdas/shared/src/common/models/utils/validation_utils.py

@@ -1,5 +1,7 @@
 """Utils for backend folder"""
 
+from typing import Any
+
 from fhir.resources.R4B.identifier import Identifier
 
 from common.models.constants import RedisHashKeys, Urls
@@ -61,11 +63,14 @@ def convert_disease_codes_to_vaccine_type(
     return vaccine_type
 
 
-def get_vaccine_type(immunization: dict):
+def get_vaccine_type(immunization: dict | Any) -> str:
     """
-    Take a FHIR immunization resource and returns the vaccine type based on the combination of target diseases.
-    If combination of disease types does not map to a valid vaccine type, a value error is raised
+    Take a FHIR immunization resource (dict or Immunization model) and returns the vaccine type
+    based on the combination of target diseases. If combination of disease types does not map
+    to a valid vaccine type, a value error is raised.
     """
+    if not isinstance(immunization, dict):
+        immunization = immunization.dict()
     # Obtain list of target diseases
     try:
         target_diseases = get_target_disease_codes(immunization)