#51

CederGroupHub · Feb 25, 2024 · a363ad5 · a363ad5
1 parent cbe680b
commit a363ad5
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 12 deletions.
diff --git a/alab_management/alarm.py b/alab_management/alarm.py
@@ -4,6 +4,7 @@
 
 from slack_sdk import WebClient
 from slack_sdk.errors import SlackApiError
+from .config import AlabOSConfig
 
 
 def format_message_to_codeblock(message: str) -> str:
@@ -71,6 +72,7 @@ def __init__(
             slack_bot_token: The slack bot token to send the alert from.
             slack_channel_id: The slack channel id to send the alert to.
         """
+        self.sim_mode_flag = AlabOSConfig().is_sim_mode()
         self.email_alert = False
         self.slack_alert = False
         if email_receivers is not None:
@@ -132,10 +134,7 @@ def alert(self, message: str, category: str):
             category: The category of the message.
         """
         import os
-        sim_mode_flag = os.getenv("SIM_MODE_FLAG", "True")
-        sim_mode_flag_boolean = sim_mode_flag.lower() == "true"
-        sim_mode_flag = sim_mode_flag_boolean
-        if not sim_mode_flag:
+        if not self.sim_mode_flag:
             try:
                 if self.platforms["email"]:
                     self.send_email(message, category)

diff --git a/alab_management/device_manager.py b/alab_management/device_manager.py
@@ -114,8 +114,8 @@ def __init__(self, _check_status: bool = True):
               running commands. (disable it only for test purpose).
         """
         load_definition()
-        sim_mode_flag = os.getenv("SIM_MODE_FLAG", "True")
-        if sim_mode_flag.lower() == "true":
+        self.sim_mode_flag = AlabOSConfig().is_sim_mode()
+        if self.sim_mode_flag:
             self._rpc_queue_name = (
                 AlabOSConfig()["general"]["name"] + "_sim" + DEFAULT_SERVER_QUEUE_SUFFIX
             )
@@ -253,8 +253,8 @@ def __init__(self, task_id: ObjectId, timeout: int = None):
         """
         assert task_id is not None, "task_id cannot be None!"
 
-        sim_mode_flag = os.getenv("SIM_MODE_FLAG", "True")
-        if sim_mode_flag.lower() == "true":
+        self.sim_mode_flag = AlabOSConfig().is_sim_mode()
+        if self.sim_mode_flag:
             self._rpc_queue_name = (
                 AlabOSConfig()["general"]["name"] + "_sim" + DEFAULT_SERVER_QUEUE_SUFFIX
             )

diff --git a/alab_management/utils/data_objects.py b/alab_management/utils/data_objects.py
@@ -15,6 +15,7 @@
 from pymongo import collection, database
 
 from .db_lock import MongoLock
+from alab_management.config import AlabOSConfig
 
 
 class _BaseGetMongoCollection(ABC):
@@ -54,8 +55,8 @@ def init(cls):
             username=db_config.get("username", ""),
             password=db_config.get("password", ""),
         )
-        sim_mode_flag = os.getenv("SIM_MODE_FLAG", "True")
-        if sim_mode_flag.lower() == "true":
+        sim_mode_flag = AlabOSConfig().is_sim_mode()
+        if sim_mode_flag:
             cls.db = cls.client[AlabOSConfig()["general"]["name"] + "_sim"]
         else:
             cls.db = cls.client[AlabOSConfig()["general"]["name"]]  # type: ignore # pylint: disable=unsubscriptable-object
@@ -83,8 +84,8 @@ def init(cls):
             username=db_config.get("username", ""),
             password=db_config.get("password", ""),
         )
-        sim_mode_flag = os.getenv("SIM_MODE_FLAG", "True")
-        if sim_mode_flag.lower() == "true":
+        sim_mode_flag = AlabOSConfig().is_sim_mode()
+        if sim_mode_flag:
             cls.db = cls.client[AlabOSConfig()["general"]["name"] + "(completed)" + "_sim"]
         else:
             cls.db = cls.client[AlabOSConfig()["general"]["name"] + "(completed)"]