Change to use Enum

[osm/common.git] / osm_common / dataclasses / temporal_dataclasses.py

diff --git a/osm_common/dataclasses/temporal_dataclasses.py b/osm_common/dataclasses/temporal_dataclasses.py
index de04186..37736c7 100644 (file)
--- a/osm_common/dataclasses/temporal_dataclasses.py
+++ b/osm_common/dataclasses/temporal_dataclasses.py
@@ -16,7 +16,9 @@
 #######################################################################################
 
 from dataclasses import dataclass
+from enum import auto, IntEnum
 
+#######################################################################################
 # Workflow Dataclasses
 
 
@@ -37,10 +39,12 @@ class VimOperationInput:
         by the workflow at the end to update the status of the
         operation in Mongo vim collection.
     """
+
     vim_uuid: str
     op_id: str
 
 
+#######################################################################################
 # Activity Dataclasses
 
 
@@ -55,9 +59,16 @@ class TestVimConnectivityInput:
         The UUID of the VIM account as stored in the OSM vim
         collection in Mongo
     """
+
     vim_uuid: str
 
 
+class VimState(IntEnum):
+    PROCESSING = auto()
+    ENABLED = auto()
+    ERROR = auto()
+
+
 @dataclass
 class UpdateVimStateInput:
     """
@@ -69,7 +80,7 @@ class UpdateVimStateInput:
         The UUID of the VIM account as stored in the OSM vim
         collection in Mongo
 
-    operational_state : str
+    operational_state : VimState
         A representation of the operational state (ENABLED or ERROR)
         of the VIM.
 
@@ -78,11 +89,17 @@ class UpdateVimStateInput:
         operational state, such as the error message associated
         with the ERROR operational_state.
     """
+
     vim_uuid: str
-    operational_state: str
+    operational_state: VimState
     message: str
 
 
+class VimOperationState(IntEnum):
+    COMPLETED = auto()
+    FAILED = auto()
+
+
 @dataclass
 class UpdateVimOperationStateInput:
     """
@@ -99,7 +116,7 @@ class UpdateVimOperationStateInput:
         The operation (task) id for this workflow.  This is used
         to update the status of the operation in Mongo vim collection.
 
-    op_state : str
+    op_state : VimOperationState
         A representation of the state of the specified operation id,
         such as COMPLETED, or FAILED.
 
@@ -108,7 +125,24 @@ class UpdateVimOperationStateInput:
         operation state, such as the error message explaining why
         the operation failed.
     """
+
     vim_uuid: str
     op_id: str
-    op_state: str
+    op_state: VimOperationState
     message: str
+
+
+@dataclass
+class DeleteVimInput:
+    """
+    Input dataclass for deleting vim record from the database
+
+    Attributes:
+    -----------
+    vim_uuid : str
+        The UUID of the VIM account as stored in the OSM vim
+        collection in Mongo
+
+    """
+
+    vim_uuid: str