From 334fd40afe7b2dc011b8c85f43cdf6ebba052b77 Mon Sep 17 00:00:00 2001
From: Jon Renshaw <jon.renshaw@nccgroup.com>
Date: Fri, 26 Jul 2024 15:56:56 +0100
Subject: [PATCH] Fixed json load for Classification and added test data to
 input.json

---
 pytm/pytm.py     | 9 +++++++--
 tests/input.json | 3 ++-
 2 files changed, 9 insertions(+), 3 deletions(-)

diff --git a/pytm/pytm.py b/pytm/pytm.py
index a711687..6ee4543 100644
--- a/pytm/pytm.py
+++ b/pytm/pytm.py
@@ -181,9 +181,14 @@ def __set__(self, instance, value):
 
 class varClassification(var):
     def __set__(self, instance, value):
-        if not isinstance(value, Classification):
+        if isinstance(value, str):
+            Classification[value.lstrip("Classification.")]
+
+        elif isinstance(value, Classification):
+            super().__set__(instance, value)            
+
+        elif not isinstance(value, Classification):
             raise ValueError("expecting a Classification, got a {}".format(type(value)))
-        super().__set__(instance, value)
 
 
 class varLifetime(var):
diff --git a/tests/input.json b/tests/input.json
index 0cdcff8..71344eb 100644
--- a/tests/input.json
+++ b/tests/input.json
@@ -5,7 +5,8 @@
   "onDuplicates": "IGNORE",
   "boundaries": [
     {
-      "name": "Internet"
+      "name": "Internet",
+      "maxClassification": "Classification.UNKNOWN"
     },
     {
       "name": "Server/DB"