Use enum for bidirectional mappings

deploy-agent/deployd/client/serverless_client.py

@@ -26,7 +26,7 @@
 
 log: logging.Logger = logging.getLogger(__name__)
 
-_DEPLOY_STAGE_TRANSITIONS = dict([(i, i+1) for i in range(DeployStage.PRE_DOWNLOAD, DeployStage.SERVING_BUILD)])
+_DEPLOY_STAGE_TRANSITIONS = { DeployStage(i).value : DeployStage(i+1).value for i in range(DeployStage.PRE_DOWNLOAD.value, DeployStage.SERVING_BUILD.value) }
 
 
 class ServerlessClient(BaseClient):
@@ -38,7 +38,7 @@ class ServerlessClient(BaseClient):
       PRE_DOWNLOAD->DOWNLOADING->POST_DOWNLOAD->STAGING->PRE_RESTART
           ->RESTARTING->POST_RESTART->SERVING_BUILD
     """
-    def __init__(self, env_name, stage, build, script_variables, deploy_stage: Optional[DeployStage] = None) -> None:
+    def __init__(self, env_name, stage, build, script_variables, deploy_stage: Optional[int] = None) -> None:
         """build contains build information in json format. It contains information defined in types/build.py.
         """
         self._env_name: str = utils.check_not_none(env_name, 'env_name can not be None')
@@ -47,7 +47,7 @@ def __init__(self, env_name, stage, build, script_variables, deploy_stage: Optio
 
         self._script_variables: dict[str, str] = json.loads(utils.check_not_none(script_variables, 'script_variables can not be None'))
         self._deploy_id: str = uuid.uuid4().hex
-        self._deploy_stage: DeployStage = deploy_stage if deploy_stage is not None else DeployStage.PRE_DOWNLOAD
+        self._deploy_stage: int = deploy_stage if deploy_stage is not None else DeployStage.PRE_DOWNLOAD.value
 
     def send_reports(self, env_reports=None) -> Optional[PingResponse]:
         reports: list = [status.report for status in env_reports.values()]
@@ -71,11 +71,11 @@ def _create_response(self, report) -> Optional[PingResponse]:
         if report.errorCode != 0:
             # terminate the deployment.
             return None
-        numeric_deploy_stage = DeployStage._NAMES_TO_VALUES[report.deployStage] 
+        numeric_deploy_stage = DeployStage[report.deployStage].value
         if report.status == AgentStatus.SUCCEEDED:
             # check if this is the last deploy stage.
-            if numeric_deploy_stage == DeployStage.SERVING_BUILD:
-                return PingResponse({'opCode': OperationCode.NOOP})
+            if DeployStage(numeric_deploy_stage) == DeployStage.SERVING_BUILD:
+                return PingResponse({'opCode': OperationCode.NOOP.value})
 
             # move to next deploy stage
             next_deploy_stage = _DEPLOY_STAGE_TRANSITIONS.get(numeric_deploy_stage)
@@ -88,14 +88,14 @@ def _create_response(self, report) -> Optional[PingResponse]:
         # terminate deployment
         return None
 
-    def _new_response_value(self, numeric_deploy_stage) -> PingResponse:
-        value= {'opCode': OperationCode.DEPLOY,
+    def _new_response_value(self, numeric_deploy_stage: int) -> PingResponse:
+        value= {'opCode': OperationCode.DEPLOY.value,
                 'deployGoal': {'deployId': self._deploy_id,
                                'envId': self._env_id,
                                'envName': self._env_name,
                                'stageName': self._stage,
                                'build': self._build,
                                'deployStage': numeric_deploy_stage}}
-        if numeric_deploy_stage == DeployStage.PRE_DOWNLOAD:
+        if DeployStage(numeric_deploy_stage) == DeployStage.PRE_DOWNLOAD:
             value['deployGoal']['scriptVariables'] = self._script_variables
         return PingResponse(jsonValue=value)

deploy-agent/deployd/common/config.py

@@ -88,12 +88,12 @@ def update_variables(self, deploy_status):
         # TODO: This is only used for migration, should clean them up
         if isinstance(deploy_status.report.deployStage, int):
             self._environ['DEPLOY_STEP'] = \
-                DeployStage._VALUES_TO_NAMES[deploy_status.report.deployStage]
+                DeployStage(deploy_status.report.deployStage).name
         else:
             self._environ['DEPLOY_STEP'] = deploy_status.report.deployStage
 
         if isinstance(deploy_status.op_code, int):
-            op_code = OperationCode._VALUES_TO_NAMES[deploy_status.op_code]
+            op_code = OperationCode(deploy_status.op_code).name
         else:
             op_code = deploy_status.op_code
         self._environ['OPCODE'] = op_code

deploy-agent/deployd/common/types.py

@@ -190,7 +190,7 @@ def load_from_json(self, json_value):
         self.runtime_config = json_value.get('runtime_config')
         op_code = json_value.get('op_code', OpCode.NOOP)
         if isinstance(op_code, int):
-            self.op_code = OperationCode._VALUES_TO_NAMES[op_code]
+            self.op_code = OperationCode(op_code).name
         else:
             self.op_code = op_code
 

deploy-agent/deployd/types/agent_status.py

@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import enum
 
 
-class AgentStatus(object):
+class AgentStatus(enum.Enum):
     # Deploy step successfully completed
     SUCCEEDED = 0
     # Deploy step is in unknown status, mostly set in the beginning of the step
@@ -35,29 +36,3 @@ class AgentStatus(object):
     RUNTIME_MISMATCH = 8
     # Use for deploy-agent internally
     ABORTED_BY_SERVER = 9
-
-    _VALUES_TO_NAMES = {
-        0: "SUCCEEDED",
-        1: "UNKNOWN",
-        2: "AGENT_FAILED",
-        3: "RETRYABLE_AGENT_FAILED",
-        4: "SCRIPT_FAILED",
-        5: "ABORTED_BY_SERVICE",
-        6: "SCRIPT_TIMEOUT",
-        7: "TOO_MANY_RETRY",
-        8: "RUNTIME_MISMATCH",
-        9: "ABORTED_BY_SERVER",
-    }
-
-    _NAMES_TO_VALUES = {
-        "SUCCEEDED": 0,
-        "UNKNOWN": 1,
-        "AGENT_FAILED": 2,
-        "RETRYABLE_AGENT_FAILED": 3,
-        "SCRIPT_FAILED": 4,
-        "ABORTED_BY_SERVICE": 5,
-        "SCRIPT_TIMEOUT": 6,
-        "TOO_MANY_RETRY": 7,
-        "RUNTIME_MISMATCH": 8,
-        "ABORTED_BY_SERVER": 9,
-    }

deploy-agent/deployd/types/deploy_goal.py

@@ -38,7 +38,8 @@ def __init__(self, jsonValue=None):
             self.stageType = jsonValue.get('stageType')
             # TODO: Only used for migration, should remove later
             if isinstance(jsonValue.get('deployStage'), int):
-                self.deployStage = DeployStage._VALUES_TO_NAMES[jsonValue.get('deployStage')]
+                self.deployStage = DeployStage(
+                    jsonValue.get('deployStage')).name
             else:
                 self.deployStage = jsonValue.get('deployStage')
 

deploy-agent/deployd/types/deploy_stage.py

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import enum
 
-class DeployStage(object):
+
+class DeployStage(enum.Enum):
     # Reserved by deploy system
     UNKNOWN = 0
     # (Optional) Unregister with LB/ZK, drain traffic, shut down service etc.
@@ -36,31 +38,3 @@ class DeployStage(object):
     STOPPING = 9
     # Complete shutting down the service
     STOPPED = 10
-
-    _VALUES_TO_NAMES = {
-        0: "UNKNOWN",
-        1: "PRE_DOWNLOAD",
-        2: "DOWNLOADING",
-        3: "POST_DOWNLOAD",
-        4: "STAGING",
-        5: "PRE_RESTART",
-        6: "RESTARTING",
-        7: "POST_RESTART",
-        8: "SERVING_BUILD",
-        9: "STOPPING",
-        10: "STOPPED",
-    }
-
-    _NAMES_TO_VALUES = {
-        "UNKNOWN": 0,
-        "PRE_DOWNLOAD": 1,
-        "DOWNLOADING": 2,
-        "POST_DOWNLOAD": 3,
-        "STAGING": 4,
-        "PRE_RESTART": 5,
-        "RESTARTING": 6,
-        "POST_RESTART": 7,
-        "SERVING_BUILD": 8,
-        "STOPPING": 9,
-        "STOPPED": 10,
-    }

deploy-agent/deployd/types/deploy_type.py

@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import enum
 
 
-class DeployType(object):
+class DeployType(enum.Enum):
     # Regular deploy
     REGULAR = 0
     # Special deploy which should go fast
@@ -24,19 +25,3 @@ class DeployType(object):
     RESTART = 3
     # Special deploy to stop service
     STOP = 4
-
-    _VALUES_TO_NAMES = {
-        0: "REGULAR",
-        1: "HOTFIX",
-        2: "ROLLBACK",
-        3: "RESTART",
-        4: "STOP",
-    }
-
-    _NAMES_TO_VALUES = {
-        "REGULAR": 0,
-        "HOTFIX": 1,
-        "ROLLBACK": 2,
-        "RESTART": 3,
-        "STOP": 4,
-    }

deploy-agent/deployd/types/opcode.py

@@ -11,9 +11,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
+import enum
 
 
-class OperationCode(object):
+class OperationCode(enum.Enum):
     # No action needed
     NOOP = 0
     # Agent need to install new build
@@ -33,27 +34,3 @@ class OperationCode(object):
     ROLLBACK = 7
     # Agent needs to shutdown service
     STOP = 8
-
-    _VALUES_TO_NAMES = {
-        0: "NOOP",
-        1: "DEPLOY",
-        2: "UPDATE",
-        3: "RESTART",
-        4: "DELETE",
-        5: "TERMINATE",
-        6: "WAIT",
-        7: "ROLLBACK",
-        8: "STOP",
-    }
-
-    _NAMES_TO_VALUES = {
-        "NOOP": 0,
-        "DEPLOY": 1,
-        "UPDATE": 2,
-        "RESTART": 3,
-        "DELETE": 4,
-        "TERMINATE": 5,
-        "WAIT": 6,
-        "ROLLBACK": 7,
-        "STOP": 8,
-    }

deploy-agent/deployd/types/ping_report.py

@@ -40,12 +40,14 @@ def __init__(self, jsonValue=None):
             self.deployId = jsonValue.get('deployId')
             self.envId = jsonValue.get('envId')
             if isinstance(jsonValue.get('deployStage'), int):
-                self.deployStage = DeployStage._VALUES_TO_NAMES[jsonValue.get('deployStage')]
+                self.deployStage = DeployStage(
+                    jsonValue.get('deployStage')).name
             else:
                 self.deployStage = jsonValue.get('deployStage')
 
             if isinstance(jsonValue.get('status'), int):
-                self.status = AgentStatus._VALUES_TO_NAMES[jsonValue.get('status')]
+                self.status = AgentStatus(jsonValue.get(
+                    'status')).name
             else:
                 self.status = jsonValue.get('status')
 

deploy-agent/deployd/types/ping_request.py

@@ -61,12 +61,13 @@ def to_json(self):
 
             # TODO: Only used for migration, should remove later
             if isinstance(report.deployStage, int):
-                ping_report["deployStage"] = DeployStage._VALUES_TO_NAMES[report.deployStage]
+                ping_report["deployStage"] = DeployStage(
+                    report.deployStage).name
             else:
                 ping_report["deployStage"] = report.deployStage
 
             if isinstance(report.status, int):
-                ping_report["agentStatus"] = AgentStatus._VALUES_TO_NAMES[report.status]
+                ping_report["agentStatus"] = AgentStatus(report.status).name
             else:
                 ping_report["agentStatus"] = report.status
 

deploy-agent/deployd/types/ping_response.py

@@ -25,7 +25,7 @@ def __init__(self, jsonValue=None):
         if jsonValue:
             # TODO: Only used for migration, should remove later
             if isinstance(jsonValue.get('opCode'), int):
-                self.opCode = OperationCode._VALUES_TO_NAMES[jsonValue.get('opCode')]
+                self.opCode = OperationCode(jsonValue.get('opCode')).name
             else:
                 self.opCode = jsonValue.get('opCode')
 

deploy-agent/tests/unit/deploy/client/test_serverless_client.py

@@ -16,10 +16,11 @@
 import unittest
 from tests import TestCase
 
-from deployd.client.serverless_client import ServerlessClient
-from deployd.common.types import DeployStage, DeployStatus, AgentStatus
+from deployd.client.serverless_client import ServerlessClient, _DEPLOY_STAGE_TRANSITIONS
+from deployd.common.types import DeployStatus, AgentStatus
 from deployd.types.ping_report import PingReport
 from deployd.types.ping_response import PingResponse
+from deployd.types.deploy_stage import DeployStage
 
 
 class TestServerlessClient(TestCase):
@@ -64,7 +65,7 @@ def test_deploy_stage_transition(self) -> None:
 
     def test_run_with_defined_deploy_stage(self) -> None:
         self.client = ServerlessClient(env_name=self.env_name, stage=self.stage, build=self.build,
-                                       script_variables=self.script_variables, deploy_stage=DeployStage.PRE_RESTART)
+                                       script_variables=self.script_variables, deploy_stage=DeployStage.PRE_RESTART.value)
         report: PingReport = self._new_report()
         report.deployId = None
         deploy_status = DeployStatus()
@@ -103,6 +104,10 @@ def test_unknow_status_cause_retry(self):
 
         response = self.client.send_reports(env_status)
         self.assertEqual(response.deployGoal.deployStage, 'PRE_DOWNLOAD')
+
+    def test_deploy_stage_transitions(self):
+        expected: dict[int, int] = dict([(i, i+1) for i in range(DeployStage.PRE_DOWNLOAD.value, DeployStage.SERVING_BUILD.value)])
+        self.assertDictEqual(_DEPLOY_STAGE_TRANSITIONS, expected)
 
     def test_create_response_last_deploy_stage(self):
         report = self._new_report()