amend function

uktrade · Jul 3, 2024 · 4a49311 · 4a49311
1 parent d619c0e
commit 4a49311
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/mail/libraries/lite_to_edifact_converter.py b/mail/libraries/lite_to_edifact_converter.py
@@ -137,11 +137,11 @@ def generate_lines_for_licence(licence: LicencePayload) -> Iterable[chieftypes._
                 controlled_by = "Q"  # usage is controlled by quantity only
                 quantity = commodity.get("quantity")
 
-                try:
-                    qunit = UnitMapping[commodity["unit"]]
-                except KeyError:
+                if commodity["unit"] in LegacyUnitCodeMapping:
                     qunit = LegacyUnitCodeMapping[commodity["unit"]]
-
+                else:
+                    qunit = UnitMapping[commodity["unit"]]
+
                 if qunit == UnitMapping.NAR:
                     quantity = int(quantity)
                 yield chieftypes.LicenceDataLine(