Merge pull request #97 from uktrade/fix-country-id-territory-code

stripping territory code while building the data file
uktrade · Feb 16, 2022 · e0c0e2e · e0c0e2e
2 parents 7a6474e + 06569e6
commit e0c0e2e
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/mail/libraries/helpers.py b/mail/libraries/helpers.py
@@ -334,11 +334,13 @@ def get_action(reference) -> str:
 def get_country_id(country):
     try:
         if type(country) == dict:
-            return country["id"]
+            country_code = country["id"]
         else:
-            return json.loads(country)["id"]
+            country_code = json.loads(country)["id"]
+        # skip territory code, if exists
+        return country_code.split("-")[0]
     except (TypeError, JSONDecodeError):
-        return country
+        return country.split("-")[0]
 
 
 def sort_dtos_by_date(input_dtos):

diff --git a/mail/libraries/lite_to_edifact_converter.py b/mail/libraries/lite_to_edifact_converter.py
@@ -103,7 +103,8 @@ def licences_to_edifact(licences: QuerySet, run_number: int) -> str:
             elif licence_payload.get("type") in LicenceTypeEnum.STANDARD_LICENCES:
                 if licence_payload.get("end_user"):
                     # In the absence of country_group or countries use country of End user
-                    country_id = licence_payload.get("end_user").get("address").get("country").get("id")
+                    country = licence_payload.get("end_user").get("address").get("country")
+                    country_id = get_country_id(country)
                     line_no += 1
                     edifact_file += "\n{}\\country\\{}\\\\{}".format(line_no, country_id, "D")
 

diff --git a/mail/tests/test_helpers.py b/mail/tests/test_helpers.py
@@ -94,7 +94,9 @@ def test_transaction_reference_for_licence_data(self, reference, transaction_ref
     def test_action_reference_for_usage(self, reference, action):
         self.assertEqual(get_action(reference), action)
 
-    @parameterized.expand([("GB", "GB"), ('{"id": "GB"}', "GB"), ({"id": "GB"}, "GB")])
+    @parameterized.expand(
+        [("GB", "GB"), ('{"id": "GB"}', "GB"), ({"id": "GB"}, "GB"), ("AE-DU", "AE"), ({"id": "AE-DU"}, "AE"),]
+    )
     @tag("country-id")
     def test_get_country_id(self, country, id):
         self.assertEqual(get_country_id(country), id)