VED-79: fix idempotency issues

Akol125 · Akol125 · commit da3d95856a93 · 2025-07-29T13:05:42.000+01:00
diff --git a/mns_subscription/models/errors.py b/mns_subscription/models/errors.py
@@ -11,8 +11,8 @@ class Severity(str, Enum):
 class Code(str, Enum):
     forbidden = "forbidden"
     not_found = "not-found"
-    invalid = "invalid"
-    server_error = "exception"
+    invalid = "invalid or missing access token"
+    server_error = "internal server error"
     invariant = "invariant"
     not_supported = "not-supported"
     duplicate = "duplicate"
@@ -34,21 +34,34 @@ def to_operation_outcome() -> dict:
 
 
 @dataclass
-class UnauthorizedVaxOnRecordError(RuntimeError):
+class TokenValidationError(RuntimeError):
     @staticmethod
     def to_operation_outcome() -> dict:
-        msg = "Unauthorized request for vaccine type present in the stored immunization resource"
+        msg = "Missing/Invalid Token"
         return create_operation_outcome(
             resource_id=str(uuid.uuid4()),
             severity=Severity.error,
-            code=Code.forbidden,
+            code=Code.invalid,
+            diagnostics=msg,
+        )
+
+
+@dataclass
+class ConflictError(RuntimeError):
+    @staticmethod
+    def to_operation_outcome() -> dict:
+        msg = "Conflict"
+        return create_operation_outcome(
+            resource_id=str(uuid.uuid4()),
+            severity=Severity.error,
+            code=Code.duplicate,
             diagnostics=msg,
         )
 
 
 @dataclass
 class ResourceFoundError(RuntimeError):
-    """Return this error when the requested FHIR resource does exist"""
+    """Return this error when the requested resource does not exist or not complete"""
 
     resource_type: str
     resource_id: str
@@ -67,7 +80,7 @@ def to_operation_outcome(self) -> dict:
 
 @dataclass
 class UnhandledResponseError(RuntimeError):
-    """Use this error when the response from an external service (ex: dynamodb) can't be handled"""
+    """Use this unhandled errors"""
 
     response: dict | str
     message: str
@@ -85,21 +98,21 @@ def to_operation_outcome(self) -> dict:
 
 
 @dataclass
-class IdentifierDuplicationError(RuntimeError):
-    """Fine grain validation"""
+class ServerError(RuntimeError):
+    """Use when there is a server error"""
 
-    identifier: str
+    response: dict | str
+    message: str
 
-    def __str__(self) -> str:
-        return f"The provided identifier: {self.identifier} is duplicated"
+    def __str__(self):
+        return f"{self.message}\n{self.response}"
 
     def to_operation_outcome(self) -> dict:
-        msg = self.__str__()
         return create_operation_outcome(
             resource_id=str(uuid.uuid4()),
             severity=Severity.error,
-            code=Code.duplicate,
-            diagnostics=msg,
+            code=Code.server_error,
+            diagnostics=self.__str__(),
         )
 
 
diff --git a/mns_subscription/src/mns_service.py b/mns_subscription/src/mns_service.py
@@ -4,7 +4,13 @@
 import logging
 import json
 from authentication import AppRestrictedAuth
-from models.errors import UnhandledResponseError
+from models.errors import (
+    UnhandledResponseError,
+    ResourceFoundError,
+    UnauthorizedError,
+    ServerError,
+    TokenValidationError
+)
 
 SQS_ARN = os.getenv("SQS_ARN")
 MNS_URL = "https://int.api.service.nhs.uk/multicast-notification-service/subscriptions"
@@ -13,39 +19,102 @@
 class MnsService:
     def __init__(self, authenticator: AppRestrictedAuth):
         self.authenticator = authenticator
-
-        logging.info(f"Using SQS ARN for subscription: {SQS_ARN}")
-
-    def subscribe_notification(self) -> dict | None:
-        access_token = self.authenticator.get_access_token()
-        request_headers = {
+        self.access_token = self.authenticator.get_access_token()
+        self.request_headers = {
             'Content-Type': 'application/fhir+json',
-            'Authorization': f'Bearer {access_token}',
+            'Authorization': f'Bearer {self.access_token}',
             'X-Correlation-ID': str(uuid.uuid4())
         }
-
-        subscription_payload = {
+        self.subscription_payload = {
             "resourceType": "Subscription",
             "status": "requested",
-            "reason": "Subscribe SQS to MNS test-signal",
+            "reason": "Subscribe SQS to NHS Number Change Events",
             "criteria": "eventType=nhs-number-change-2",
             "channel": {
                 "type": "message",
                 "endpoint": SQS_ARN,
                 "payload": "application/json"
                 }
             }
-        response = requests.post(MNS_URL, headers=request_headers, data=json.dumps(subscription_payload))
 
-        print(f"Access Token: {access_token}")
+        logging.info(f"Using SQS ARN for subscription: {SQS_ARN}")
+
+    def subscribe_notification(self) -> dict | None:
+
+        response = requests.post(MNS_URL, headers=self.request_headers, data=json.dumps(self.subscription_payload))
+
+        print(f"Access Token: {self.access_token}")
         print(f"SQS ARN: {SQS_ARN}")
-        print(f"Headers: {request_headers}")
-        print(f"Payload: {json.dumps(subscription_payload, indent=2)}")
+        print(f"Headers: {self.request_headers}")
+        print(f"Payload: {json.dumps(self.subscription_payload, indent=2)}")
 
-        if response.status_code == 201:
+        if response.status_code == 200:
             return response.json()
+        elif response.status_code == 409:
+            msg = "SQS Queue Already Subscribed, can't re-subscribe"
+            raise UnhandledResponseError(response=response.json(), message=msg)
+        elif response.status_code == 401:
+            msg = "SQS Queue Already Subscribed, can't re-subscribe"
+            raise TokenValidationError(response=response.json(), message=msg)
+        elif response.status_code == 400:
+            msg = "Resource Type provided for this is not correct"
+            raise ResourceFoundError(response=response.json(), message=msg)
+        elif response.status_code == 403:
+            msg = "You don't have the right permissions for this request"
+            raise UnauthorizedError(response=response.json(), message=msg)
+        elif response.status_code == 500:
+            msg = "Internal Server Error"
+            raise ServerError(response=response.json(), message=msg)
+        else:
+            msg = f"Unhandled error: {response.status_code} - {response.text}"
+            raise UnhandledResponseError(response=response.json(), message=msg)
+
+    def get_subscription(self) -> dict | None:
+        response = requests.get(MNS_URL, headers=self.request_headers)
+        logging.info(f"GET {MNS_URL}")
+        logging.debug(f"Headers: {self.request_headers}")
+
+        if response.status_code == 200:
+            bundle = response.json()
+            # Assume a FHIR Bundle with 'entry' list
+            for entry in bundle.get("entry", []):
+                resource = entry.get("channel", {})
+                channel = resource.get("channel", {})
+                if channel.get("endpoint") == SQS_ARN:
+                    return resource  # Found a matching subscription
+            return None  # No subscription for this SQS ARN
         elif response.status_code == 404:
             return None
+        elif response.status_code == 401:
+            msg = "Token validation failed for the request"
+            raise TokenValidationError(response=response.json(), message=msg)
+        elif response.status_code == 400:
+            msg = "Bad request: Resource type or parameters incorrect"
+            raise ResourceFoundError(response=response.json(), message=msg)
+        elif response.status_code == 403:
+            msg = "You don't have the right permissions for this request"
+            raise UnauthorizedError(response=response.json(), message=msg)
+        elif response.status_code == 500:
+            msg = "Internal Server Error"
+            raise ServerError(response=response.json(), message=msg)
         else:
-            msg = "MNS subscription failed"
+            msg = f"Unhandled error: {response.status_code} - {response.text}"
             raise UnhandledResponseError(response=response.json(), message=msg)
+
+    def check_subscription(self) -> dict:
+        """
+        Ensures that a subscription exists for this SQS_ARN.
+        If not found, creates one.
+        Returns the subscription.
+        """
+        try:
+            existing = self.get_subscription()
+            if existing:
+                logging.info("Subscription for this SQS ARN already exists.")
+                return existing
+            else:
+                logging.info("No subscription found for this SQS ARN. Creating new subscription...")
+                return self.subscribe_notification()
+        except Exception as e:
+            logging.error(f"Error ensuring subscription: {e}")
+            raise
diff --git a/mns_subscription/src/subscribe_mns.py b/mns_subscription/src/subscribe_mns.py
@@ -26,7 +26,7 @@ def run_subscription():
         mns = MnsService(authenticator)
 
         logging.info("Subscribing to MNS...")
-        result = mns.subscribe_notification()
+        result = mns.check_subscription()
         logging.info(f"Subscription Result: {result}")
         return result
     except Exception:
