Merge pull request #52 from gamultong/hotfix/fix-chores

Payload fromdict에 enum 추가
gamultong · Dec 5, 2024 · 2771cc1 · 2771cc1
2 parents 6e14768 + b418cd5
commit 2771cc1
Show file tree

Hide file tree

Showing 3 changed files with 45 additions and 2 deletions.
diff --git a/message/payload/internal/base_payload.py b/message/payload/internal/base_payload.py
@@ -1,6 +1,8 @@
 from .exceptions import InvalidFieldException, MissingFieldException, DumbHumanException
 from .parsable_payload import ParsablePayload
 
+from enum import Enum
+
 
 class Payload():
     @classmethod
@@ -36,6 +38,10 @@ def _from_dict(cls, dict: dict):
                     raise MissingFieldException(key, e)
                 continue
 
+            if issubclass(t, Enum) and dict[key] in t:
+                kwargs[key] = t(dict[key])
+                continue
+
             if not type(dict[key]) == t:
                 # 잘못된 형식의 data
                 raise InvalidFieldException(key, dict[key])

diff --git a/message/payload/test/testdata/base_payload_testdata.py b/message/payload/test/testdata/base_payload_testdata.py
@@ -1,11 +1,17 @@
 from message.payload import Payload, InvalidFieldException, MissingFieldException
 from dataclasses import dataclass
+from enum import Enum
+
+
+class ExampleEnum(str, Enum):
+    Hi = "hi"
 
 
 @dataclass
 class ExamplePayload(Payload):
     foo: str
     bar: str
+    enum: ExampleEnum
 
 
 @dataclass
@@ -18,23 +24,33 @@ class ExampleWrapperPayload(Payload):
 EXAMPLE_PAYLOAD_DICT = {
     "foo": "foo",
     "bar": "bar",
+    "enum": "hi"
 }
 
 EXAMPLE_PAYLOAD_DICT_MISSING_KEY = {
-    "bar": "bar"
+    "bar": "bar",
+    "enum": "hi"
 }
 
 EXAMPLE_PAYLOAD_DICT_INVALID_KEY = {
     "foo": "foo",
     "bar": "bar",
+    "enum": "hi",
     "baz": "baz"
 }
 
 EXAMPLE_PAYLOAD_DICT_INVALID_VALUE_TYPE = {
     "foo": "foo",
+    "enum": "hi",
     "bar": 1,
 }
 
+EXAMPLE_PAYLOAD_DICT_INVALID_ENUM = {
+    "foo": "foo",
+    "bar": "bar",
+    "enum": "hey",
+}
+
 EXAMPLE_PAYLOAD_WRAPPED_DICT = {
     "a": 1,
     "b": "b",
@@ -59,6 +75,12 @@ class ExampleWrapperPayload(Payload):
     "wrapped": EXAMPLE_PAYLOAD_DICT_INVALID_VALUE_TYPE
 }
 
+EXAMPLE_PAYLOAD_WRAPPED_DICT_INVALID_ENUM = {
+    "a": 1,
+    "b": "b",
+    "wrapped": EXAMPLE_PAYLOAD_DICT_INVALID_ENUM
+}
+
 EXCEPTION_TEST_CASES = \
     [
         {
@@ -73,6 +95,12 @@ class ExampleWrapperPayload(Payload):
             "exp": InvalidFieldException,
             "msg": "invaild field: bar -> 1"
         },
+        {
+            "dict": EXAMPLE_PAYLOAD_DICT_INVALID_ENUM,
+            "payload": ExamplePayload,
+            "exp": InvalidFieldException,
+            "msg": "invaild field: enum -> hey"
+        },
         {
             "dict": EXAMPLE_PAYLOAD_DICT_INVALID_KEY,
             "payload": ExamplePayload,
@@ -91,6 +119,12 @@ class ExampleWrapperPayload(Payload):
             "exp": InvalidFieldException,
             "msg": "invaild field: wrapped.bar -> 1"
         },
+        {
+            "dict": EXAMPLE_PAYLOAD_WRAPPED_DICT_INVALID_ENUM,
+            "payload": ExampleWrapperPayload,
+            "exp": InvalidFieldException,
+            "msg": "invaild field: wrapped.enum -> hey"
+        },
         {
             "dict": EXAMPLE_PAYLOAD_WRAPPED_DICT_INVALID_KEY,
             "payload": ExampleWrapperPayload,

diff --git a/server.py b/server.py
@@ -31,7 +31,10 @@ async def session(ws: WebSocket):
             msg.header = {"sender": conn.id}
             await ConnectionManager.handle_message(msg)
         except Exception as e:
-            print(f"WebSocket connection closed: {e}")
+            msg = e
+            if hasattr(e, "msg"):
+                msg = e.msg
+            print(f"WebSocket connection closed: {type(e)} '{msg}'")
             break
 
     await ConnectionManager.close(conn)