增强配置系统

2026-03-10 16:13:10 +08:00
parent a638e4203b
commit 291aa0b4b7
10 changed files with 481 additions and 45 deletions
@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+import pytest
+from pydantic import BaseModel, ConfigDict
+
+from sync_state_machine.common.registry import DomainRegistry
+from sync_state_machine.common.sync_node import SyncNode
+from sync_state_machine.config import (
+    BaseDataSourceConfig,
+    DbDataSourceConfig,
+    LoggingConfig,
+    PersistConfig,
+    PipelineRunConfig,
+    build_config,
+    register_datasource_type,
+)
+from sync_state_machine.datasource import BaseDataSource
+from sync_state_machine.datasource.db import BaseDbHandler, DbDataSource
+from sync_state_machine.pipeline.factory import create_pipeline_from_config
+from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
+
+
+class MemoryConfig(BaseDataSourceConfig):
+    model_config = ConfigDict(extra="forbid", validate_assignment=True)
+
+    type: str = "memory"
+    namespace: str = "default"
+
+
+class MemorySchema(BaseModel):
+    __test__ = False
+
+    id: str
+    name: str | None = None
+
+
+class MemoryNode(SyncNode[MemorySchema]):
+    __test__ = False
+
+    node_type = "memory_demo"
+    schema = MemorySchema
+
+
+class MemoryStrategy(DefaultSyncStrategy[MemorySchema]):
+    __test__ = False
+
+    schema = MemorySchema
+
+
+class MemoryDataSource(BaseDataSource):
+    def __init__(self, namespace: str):
+        super().__init__()
+        self.namespace = namespace
+
+
+class MemoryHandler(BaseDbHandler[MemorySchema]):
+    __test__ = False
+
+    def __init__(self, datasource: MemoryDataSource):
+        super().__init__(datasource, "memory_demo", MemoryNode, MemorySchema)
+
+    def extract_id(self, data: dict[str, str]) -> str:
+        return str(data.get("id", ""))
+
+    async def load(self):
+        return []
+
+    async def create_all(self, nodes):
+        return []
+
+    async def update_all(self, nodes):
+        return []
+
+
+class MemoryRemoteDbHandler(BaseDbHandler[MemorySchema]):
+    __test__ = False
+
+    def __init__(self, datasource: DbDataSource):
+        super().__init__(datasource, "memory_demo", MemoryNode, MemorySchema)
+
+    def extract_id(self, data: dict[str, str]) -> str:
+        return str(data.get("id", ""))
+
+    async def load(self):
+        return []
+
+    async def create_all(self, nodes):
+        return []
+
+    async def update_all(self, nodes):
+        return []
+
+
+def _ensure_memory_domain_registered() -> None:
+    if DomainRegistry.is_registered("memory_demo"):
+        reg = DomainRegistry.get_registration("memory_demo")
+        reg.schema = MemorySchema
+        reg.node_class = MemoryNode
+        reg.strategy_class = MemoryStrategy
+        reg.db_handler_class = MemoryRemoteDbHandler
+        reg.handler_classes["db"] = MemoryRemoteDbHandler
+        reg.handler_classes["memory"] = MemoryHandler
+        return
+
+    DomainRegistry.register(
+        node_type="memory_demo",
+        schema=MemorySchema,
+        node_class=MemoryNode,
+        strategy_class=MemoryStrategy,
+        db_handler_class=MemoryRemoteDbHandler,
+    )
+    DomainRegistry.register_handler("memory_demo", "memory", MemoryHandler)
+
+
+def _memory_factory(config: MemoryConfig, endpoint_name: str) -> MemoryDataSource:
+    return MemoryDataSource(namespace=f"{endpoint_name}:{config.namespace}")
+
+
+@pytest.mark.asyncio
+async def test_create_pipeline_from_config_supports_registered_custom_datasource_type(tmp_path: Path) -> None:
+    register_datasource_type("memory", config_model=MemoryConfig, factory=_memory_factory, replace=True)
+    _ensure_memory_domain_registered()
+
+    config = PipelineRunConfig(
+        node_types=["memory_demo"],
+        local_datasource=MemoryConfig(type="memory", namespace="alpha"),
+        remote_datasource=DbDataSourceConfig(type="db"),
+        strategies=[],
+        persist=PersistConfig(db_path=str(tmp_path / "memory-registry.db")),
+        logging=LoggingConfig(initialize=False),
+    )
+
+    pipeline = await create_pipeline_from_config(config)
+    try:
+        assert isinstance(pipeline.local_datasource, MemoryDataSource)
+        assert pipeline.local_datasource.namespace == "local:alpha"
+        assert isinstance(pipeline.local_datasource.get_handler("memory_demo"), MemoryHandler)
+        assert isinstance(pipeline.remote_datasource.get_handler("memory_demo"), MemoryRemoteDbHandler)
+    finally:
+        await pipeline.close()
+
+
+def test_build_config_supports_same_file_datasource_bootstrap(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
+    bootstrap_module = tmp_path / "memory_bootstrap_demo.py"
+    bootstrap_module.write_text(
+        "from pydantic import ConfigDict\n"
+        "from sync_state_machine.config import BaseDataSourceConfig, register_datasource_type\n"
+        "from sync_state_machine.datasource import BaseDataSource\n"
+        "\n"
+        "class BootstrapMemoryConfig(BaseDataSourceConfig):\n"
+        "    model_config = ConfigDict(extra=\"forbid\", validate_assignment=True)\n"
+        "    type: str = \"memory_bootstrap\"\n"
+        "    namespace: str = \"demo\"\n"
+        "\n"
+        "class BootstrapMemoryDataSource(BaseDataSource):\n"
+        "    def __init__(self, namespace: str):\n"
+        "        super().__init__()\n"
+        "        self.namespace = namespace\n"
+        "\n"
+        "def build_memory_datasource(config, endpoint_name):\n"
+        "    return BootstrapMemoryDataSource(namespace=f\"{endpoint_name}:{config.namespace}\")\n"
+        "\n"
+        "def register():\n"
+        "    register_datasource_type(\n"
+        "        \"memory_bootstrap\",\n"
+        "        config_model=BootstrapMemoryConfig,\n"
+        "        factory=build_memory_datasource,\n"
+        "        replace=True,\n"
+        "    )\n",
+        encoding="utf-8",
+    )
+    monkeypatch.syspath_prepend(str(tmp_path))
+
+    profile_path = tmp_path / "custom-memory.yaml"
+    profile_path.write_text(
+        "extensions:\n"
+        "  datasource_bootstraps:\n"
+        "    - memory_bootstrap_demo:register\n"
+        "node_types: []\n"
+        "local_datasource:\n"
+        "  type: memory_bootstrap\n"
+        "  namespace: local-demo\n"
+        "remote_datasource:\n"
+        "  type: db\n"
+        "persist:\n"
+        "  db_path: ${PROJECT_ROOT}/runtime/test-memory-bootstrap.db\n"
+        "logging:\n"
+        "  initialize: false\n",
+        encoding="utf-8",
+    )
+
+    config = build_config(project_root=tmp_path, overrides_file=profile_path)
+
+    assert config.local_datasource.type == "memory_bootstrap"
+    assert getattr(config.local_datasource, "namespace") == "local-demo"
+    assert config.remote_datasource.type == "db"