清理代码

tests/unit/test_factory_endpoint_initialization.py

@@ -14,7 +14,8 @@ from sync_state_machine.config import (
     PipelineRunConfig,
     register_datasource_type,
 )
-from sync_state_machine.datasource.db import BaseDbHandler, DbDataSource
+from sync_state_machine.datasource import BaseDataSource
+from sync_state_machine.datasource.handler import BaseNodeHandler
 from sync_state_machine.pipeline.factory import create_pipeline_from_config
 from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
 
@@ -41,11 +42,12 @@ class EndpointInitStrategy(DefaultSyncStrategy[EndpointInitSchema]):
     schema = EndpointInitSchema
 
 
-class EndpointInitDbHandler(BaseDbHandler[EndpointInitSchema]):
+class EndpointInitDbHandler(BaseNodeHandler[EndpointInitSchema]):
     __test__ = False
 
-    def __init__(self, datasource: DbDataSource):
-        super().__init__(datasource, "factory_endpoint_demo", EndpointInitNode, EndpointInitSchema)
+    def __init__(self, datasource: BaseDataSource):
+        super().__init__("factory_endpoint_demo", EndpointInitNode, EndpointInitSchema)
+        self.datasource = datasource
 
     def extract_id(self, data: dict[str, str]) -> str:
         return str(data.get("id", ""))
@@ -85,7 +87,7 @@ def _register_endpoint_init_datasource_type() -> None:
     register_datasource_type(
         "endpoint_init_db",
         config_model=EndpointInitDbConfig,
-        factory=lambda config, endpoint_name: DbDataSource(),
+        factory=lambda config, endpoint_name: BaseDataSource(),
         replace=True,
     )