🦄 refactor: message extra use dataclass

src/onestep/message.py

@@ -2,6 +2,7 @@
 import sys
 import time
 import uuid
+from dataclasses import dataclass
 from traceback import TracebackException
 from typing import Optional, Any, Union
 
@@ -14,11 +15,15 @@ def __init__(self, exc_type, exc_value, exc_traceback, **kwargs):
         self.exc_value = exc_value
 
 
+@dataclass
 class Extra:
-    def __init__(self, task_id=None, publish_time=None, failure_count=0):
-        self.task_id = task_id or str(uuid.uuid4())
-        self.publish_time = publish_time or round(time.time(), 3)
-        self.failure_count = failure_count
+    task_id: str = None
+    publish_time: float = None
+    failure_count: int = 0
+
+    def __post_init__(self):
+        self.task_id = self.task_id or str(uuid.uuid4())
+        self.publish_time = self.publish_time or round(time.time(), 3)
 
     def to_dict(self):
         return {
@@ -30,9 +35,6 @@ def to_dict(self):
     def __str__(self):
         return str(self.to_dict())
 
-    def __repr__(self):
-        return f"{self.__class__.__name__}({self.task_id}, {self.publish_time}, {self.failure_count})"
-
 
 class Message:
 

tests/test_message_extra.py

@@ -0,0 +1,44 @@
+import pytest
+from onestep.message import Extra
+
+
+@pytest.fixture
+def extra_instance():
+    # 创建一个 Extra 实例供测试使用
+    return Extra()
+
+
+def test_to_dict(extra_instance):
+    # 测试 to_dict 方法是否返回正确的字典
+    expected_dict = {
+        'task_id': extra_instance.task_id,
+        'publish_time': extra_instance.publish_time,
+        'failure_count': extra_instance.failure_count,
+    }
+    assert extra_instance.to_dict() == expected_dict
+
+
+def test_str(extra_instance):
+    # 测试 __str__ 方法是否返回正确的字符串表示
+    expected_str = str(extra_instance.to_dict())
+    assert str(extra_instance) == expected_str
+
+
+def test_custom_values():
+    # 测试使用自定义值创建 Extra 实例是否正确
+    task_id = 'custom_task_id'
+    publish_time = 123456789.0
+    failure_count = 2
+
+    extra = Extra(task_id=task_id, publish_time=publish_time, failure_count=failure_count)
+    assert extra.task_id == task_id
+    assert extra.publish_time == publish_time
+    assert extra.failure_count == failure_count
+
+
+def test_default_values():
+    # 测试使用默认值创建 Extra 实例是否正确
+    extra = Extra()
+    assert isinstance(extra.task_id, str)
+    assert isinstance(extra.publish_time, float)
+    assert extra.failure_count == 0