fix: subclass naming conflict in __init_subclass__ for TaskPoller and… (#242)

Author: jamesgetx

sdks/blue-krill/CHANGE.md

@@ -1,5 +1,9 @@
 ## Change logs
 
+### 2.1.4
+
+- Fix: `TaskPoller.__init_subclass__` 和 `CallbackHandler.__init_subclass__` 中同名子类注册覆盖的问题
+
 ### 2.1.3
 
 - Fix: EncryptField 无法正常处理以特殊 header 文本开头的内容

sdks/blue-krill/blue_krill/__init__.py

@@ -15,4 +15,4 @@
 # We undertake not to change the open source license (MIT license) applicable
 # to the current version of the project delivered to anyone in the future.
 
-__version__ = "2.1.3"
+__version__ = "2.1.4"

sdks/blue-krill/blue_krill/async_utils/poll_task.py

@@ -91,6 +91,10 @@ def __init__(self, params: Dict, metadata: PollingMetadata):
         self.metadata = metadata
 
     def __init_subclass__(cls, *args, **kwargs):
+        if cls.__name__ in cls._registered_pollers:
+            raise TypeError(
+                f"A TaskPoller subclass named '{cls.__name__}' is already registered. Poller names must be unique."
+            )
         cls._registered_pollers[cls.__name__] = cls
 
     @classmethod
@@ -204,6 +208,10 @@ class CallbackHandler(ABC):
     _registered_handlers: Dict[str, Type["CallbackHandler"]] = {}
 
     def __init_subclass__(cls, *args, **kwargs):
+        if cls.__name__ in cls._registered_handlers:
+            raise TypeError(
+                f"A CallbackHandler subclass named '{cls.__name__}' is already registered. Handler names must be unique."
+            )
         cls._registered_handlers[cls.__name__] = cls
 
     @classmethod

sdks/blue-krill/pyproject.toml

@@ -9,7 +9,7 @@ name = "blue-krill"
 description = "Tools and common packages for blueking PaaS platform."
 requires-python = ">=3.9,<3.12"
 license = "MIT"
-version = "2.1.3"
+version = "2.1.4"
 # classifieres is dynamic because we want to create Python classifiers automatically
 dynamic = ["classifiers"]
 readme = "README.md"

sdks/blue-krill/tests/async_utils/test_poll_task.py

@@ -99,6 +99,10 @@ class TimeoutedPoller(BasePoller):
             def query(self) -> PollingResult:
                 return PollingResult.done(data={"foo": "bar"})
 
+            @classmethod
+            def unregister(cls):
+                del cls._registered_pollers[cls.__name__]
+
         class TimeoutedHandler(CallbackHandler):
             result = None
             poller = None
@@ -107,6 +111,10 @@ def handle(self, result, poller):
                 TimeoutedHandler.result = result
                 TimeoutedHandler.poller = poller
 
+            @classmethod
+            def unregister(cls):
+                del cls._registered_handlers[cls.__name__]
+
         metadata = PollingMetadata(retries=0, query_started_at=time.time(), queried_count=0)
         scheduler = PollTaskScheduler(TimeoutedPoller({"param": 1}, metadata), TimeoutedHandler)
         next_metadata = scheduler.run()
@@ -120,6 +128,9 @@ def handle(self, result, poller):
         assert TimeoutedHandler.poller
         assert TimeoutedHandler.poller.params == {"param": 1}
 
+        TimeoutedPoller.unregister()
+        TimeoutedHandler.unregister()
+
     def test_exception(self):
         class ExceptionPoller(BasePoller):
             def query(self) -> PollingResult:  # type: ignore
@@ -171,12 +182,12 @@ def handle(self, result, poller):
 
 class TestTaskPoller:
     def test_continue(self):
-        class DoingPoller(BasePoller):
+        class DoingTaskPoller(BasePoller):
             def query(self) -> PollingResult:
                 return PollingResult.doing()
 
-        DoingPoller.start({}, NullResultHandler)
-        assert DoingPoller.get_async_task().subtask().apply_async.called
+        DoingTaskPoller.start({}, NullResultHandler)
+        assert DoingTaskPoller.get_async_task().subtask().apply_async.called
 
     def test_no_continue(self):
         class DonePoller(BasePoller):