Define an enum for Reboot Methods

sonic-net · Sep 18, 2024 · 0e0fc79 · 0e0fc79
1 parent 49caaca
commit 0e0fc79
Showing 1 changed file with 7 additions and 4 deletions.
diff --git a/tests/host_modules/systemd_service_test.py b/tests/host_modules/systemd_service_test.py
@@ -1,9 +1,14 @@
 import sys
 import os
 import pytest
+from enum import Enum
 from unittest import mock
 from host_modules import systemd_service
 
+class RebootMethod(Enum):
+    COLD = 1
+    HALT = 3
+
 class TestSystemdService(object):
     @mock.patch("dbus.SystemBus")
     @mock.patch("dbus.service.BusName")
@@ -98,8 +103,7 @@ def test_execute_reboot_cold(self, MockInit, MockBusName, MockSystemBus):
             res_mock.configure_mock(returncode=test_ret, stderr=test_msg)
             mock_run.return_value = res_mock
 
-            # Define reboot method and instantiate the service
-            method = 1  # Method 1 corresponds to cold reboot
+            method = RebootMethod.COLD
             systemd_service_stub = systemd_service.SystemdService(systemd_service.MOD_NAME)
 
             # Execute the reboot method
@@ -126,8 +130,7 @@ def test_execute_reboot_halt(self, MockInit, MockBusName, MockSystemBus):
             res_mock.configure_mock(returncode=test_ret, stderr=test_msg)
             mock_run.return_value = res_mock
 
-            # Define reboot method and instantiate the service
-            method = 3  # Method 3 corresponds to halt reboot
+            method = RebootMethod.HALT
             systemd_service_stub = systemd_service.SystemdService(systemd_service.MOD_NAME)
 
             # Execute the reboot method