Prompt request refactor (#31)

* Refactor to use only 1 enum * Removed unsed import * Fix isort * Adds a new Message prompt request --------- Co-authored-by: raul-marquez-csa <[email protected]>
project-chip · Nov 22, 2023 · fbae043 · fbae043
1 parent 5100c93
commit fbae043
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 28 deletions.
diff --git a/app/chip_tool/test_case.py b/app/chip_tool/test_case.py
@@ -32,7 +32,7 @@
     ManualLogUploadStep,
     ManualVerificationTestStep,
 )
-from app.user_prompt_support import OptionsSelectPromptRequest
+from app.user_prompt_support.prompt_request import MessagePromptRequest
 from app.user_prompt_support.uploaded_file_support import UploadFile
 from app.user_prompt_support.user_prompt_manager import user_prompt_manager
 from app.user_prompt_support.user_prompt_support import UserPromptSupport
@@ -256,9 +256,7 @@ async def __prompt_user_for_controller_action(self, action: str) -> None:
         """
 
         prompt = f"Please do the following action on the Controller: {action}"
-        prompt_request = OptionsSelectPromptRequest(
-            prompt=prompt, options={}, timeout=60
-        )
+        prompt_request = MessagePromptRequest(prompt=prompt, timeout=60)
         await self.send_prompt_request(prompt_request)
 
     def __handle_logs(self, logs: Any) -> None:

diff --git a/app/constants/websockets_constants.py b/app/constants/websockets_constants.py
@@ -27,6 +27,7 @@ class MessageTypeEnum(str, Enum):
     PROMPT_REQUEST = "prompt_request"
     OPTIONS_REQUEST = "options_request"
     MESSAGE_REQUEST = "message_request"
+    FILE_UPLOAD_REQUEST = "file_upload_request"
     PROMPT_RESPONSE = "prompt_response"
     TEST_UPDATE = "test_update"
     TIME_OUT_NOTIFICATION = "time_out_notification"

diff --git a/app/user_prompt_support/prompt_request.py b/app/user_prompt_support/prompt_request.py
@@ -13,43 +13,51 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 #
-from enum import Enum
 from typing import Dict, Optional
 
 from pydantic import BaseModel
 
-default_timeout_s = 60  # Seconds
-
+from app.constants.websockets_constants import MessageTypeEnum
 
-class PromptRequestType(str, Enum):
-    BASE = "base"
-    OPTIONS = "options"
-    TEXT = "text"
-    FILE = "file"
+default_timeout_s = 60  # Seconds
 
 
 class PromptRequest(BaseModel):
     prompt: Optional[str]
     timeout: int = default_timeout_s
-    __type: PromptRequestType = PromptRequestType.BASE
 
     @property
-    def type(self) -> PromptRequestType:
-        return self.__type
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.INVALID_MESSAGE
 
 
 class OptionsSelectPromptRequest(PromptRequest):
-    __type = PromptRequestType.OPTIONS
     options: Dict[str, int]
 
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.OPTIONS_REQUEST
+
 
 class TextInputPromptRequest(PromptRequest):
-    __type = PromptRequestType.TEXT
     placeholder_text: Optional[str]
     default_value: Optional[str]
     regex_pattern: Optional[str]
 
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.PROMPT_REQUEST
+
 
 class UploadFilePromptRequest(PromptRequest):
-    __type = PromptRequestType.FILE
     path: str = "api/v1/test_run_execution/file_upload/"
+
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.FILE_UPLOAD_REQUEST
+
+
+class MessagePromptRequest(PromptRequest):
+    @property
+    def messageType(self) -> MessageTypeEnum:
+        return MessageTypeEnum.MESSAGE_REQUEST
diff --git a/app/user_prompt_support/user_prompt_manager.py b/app/user_prompt_support/user_prompt_manager.py
@@ -50,16 +50,9 @@ def __init__(self, prompt: PromptRequest, message_id: int) -> None:
     def as_dictionary(self) -> Dict[MessageKeysEnum, Any]:
         prompt_dict = self.prompt.dict()
         prompt_dict[MESSAGE_ID_KEY] = self.message_id
-        message_type = MessageTypeEnum.PROMPT_REQUEST
-
-        # Check if options exist
-        if "options" in prompt_dict:
-            if prompt_dict["options"]:
-                message_type = MessageTypeEnum.OPTIONS_REQUEST
-            else:
-                message_type = MessageTypeEnum.MESSAGE_REQUEST
+
         message_dict = {
-            MessageKeysEnum.TYPE: message_type,
+            MessageKeysEnum.TYPE: self.prompt.messageType,
             MessageKeysEnum.PAYLOAD: prompt_dict,
         }
         return message_dict
@@ -227,4 +220,4 @@ def __prompt_exchange_for_message(
         )
 
 
-user_prompt_manager = UserPromptManager()
+user_prompt_manager: UserPromptManager = UserPromptManager()