From 22325ca4e36e8a734b89f596de611faea6ba3870 Mon Sep 17 00:00:00 2001
From: VirajP1002 <Viraj.Patel@ons.gov.uk>
Date: Fri, 20 Dec 2024 11:59:40 +0000
Subject: [PATCH] Update type hints and use context manager for test

---
 app/services/supplementary_data.py            |  2 +-
 app/utilities/supplementary_data_parser.py    |  2 +-
 .../parser/test_supplementary_data_parser.py  | 27 +++++++++++++------
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/app/services/supplementary_data.py b/app/services/supplementary_data.py
index 2967ea44f2..6e2aa28cd3 100644
--- a/app/services/supplementary_data.py
+++ b/app/services/supplementary_data.py
@@ -131,7 +131,7 @@ def validate_supplementary_data(
     identifier: str,
     survey_id: str,
     sds_schema_version: str | None = None,
-) -> dict:
+) -> dict[str, str | dict | int | list]:
     try:
         return validate_supplementary_data_v1(
             supplementary_data=supplementary_data,
diff --git a/app/utilities/supplementary_data_parser.py b/app/utilities/supplementary_data_parser.py
index 772f42ed5c..35e73113ae 100644
--- a/app/utilities/supplementary_data_parser.py
+++ b/app/utilities/supplementary_data_parser.py
@@ -90,7 +90,7 @@ def validate_supplementary_data_v1(
     identifier: str,
     survey_id: str,
     sds_schema_version: str | None = None,
-) -> dict:
+) -> dict[str, str | dict | int | list]:
     """Validate claims required for supplementary data"""
     supplementary_data_metadata_schema = SupplementaryDataMetadataSchema(
         unknown=INCLUDE
diff --git a/tests/app/parser/test_supplementary_data_parser.py b/tests/app/parser/test_supplementary_data_parser.py
index 8df5088a62..c5243b1c87 100644
--- a/tests/app/parser/test_supplementary_data_parser.py
+++ b/tests/app/parser/test_supplementary_data_parser.py
@@ -1,3 +1,4 @@
+from contextlib import contextmanager
 from copy import deepcopy
 
 import pytest
@@ -6,6 +7,15 @@
 from app.services.supplementary_data import validate_supplementary_data
 from app.utilities.supplementary_data_parser import validate_supplementary_data_v1
 
+
+@contextmanager
+def not_raises(exception):
+    try:
+        yield
+    except exception as validation_error:
+        raise pytest.fail(f"{validation_error} RAISED")
+
+
 SUPPLEMENTARY_DATA_PAYLOAD = {
     "dataset_id": "44f1b432-9421-49e5-bd26-e63e18a30b69",
     "survey_id": "123",
@@ -70,14 +80,15 @@ def test_invalid_supplementary_dataset_version_raises_error():
     )
 
 
-def test_valid_supplementary_dataset_version():
-    validated_payload = validate_supplementary_data_v1(
-        supplementary_data=SUPPLEMENTARY_DATA_PAYLOAD,
-        dataset_id="44f1b432-9421-49e5-bd26-e63e18a30b69",
-        identifier="12345678901",
-        survey_id="123",
-        sds_schema_version="v1",
-    )
+def test_valid_supplementary_dataset_version_does_not_raise_error():
+    with not_raises(ValidationError):
+        validated_payload = validate_supplementary_data_v1(
+            supplementary_data=SUPPLEMENTARY_DATA_PAYLOAD,
+            dataset_id="44f1b432-9421-49e5-bd26-e63e18a30b69",
+            identifier="12345678901",
+            survey_id="123",
+            sds_schema_version="v1",
+        )
 
     assert validated_payload == SUPPLEMENTARY_DATA_PAYLOAD