整理了初始化过程,现在更容易被继承了。

This commit is contained in:
strepsiades
2026-03-18 16:48:53 +08:00
parent 84b3b7652a
commit 841986740c
22 changed files with 432 additions and 277 deletions
@@ -1,28 +1,23 @@
from __future__ import annotations
from pydantic import ValidationError
import pytest
from sync_state_machine.config import (
PipelineRunConfig,
JsonlDataSourceConfig,
ApiDataSourceConfig,
DbDataSourceConfig,
PersistConfig,
LoggingConfig,
)
from sync_state_machine.pipeline.factory import _resolve_datasource_target_project_ids
def _build_config(
*,
top_level: list[str],
local_ids: list[str] | None = None,
remote_ids: list[str] | None = None,
) -> PipelineRunConfig:
return PipelineRunConfig(
node_types=["project"],
target_project_ids=top_level,
local_datasource=JsonlDataSourceConfig(
type="jsonl",
jsonl_dir="./filtered_datasource/datasource/filtered",
@@ -39,49 +34,22 @@ def _build_config(
)
def test_datasource_target_ids_take_precedence_over_top_level() -> None:
cfg = _build_config(top_level=["p_top"], local_ids=["p_local"], remote_ids=["p_remote"])
def test_pipeline_config_has_no_top_level_target_project_ids() -> None:
cfg = _build_config(local_ids=["p_local"], remote_ids=["p_remote"])
local_ids, remote_ids = _resolve_datasource_target_project_ids(cfg)
assert local_ids == ["p_local"]
assert remote_ids == ["p_remote"]
assert cfg.local_datasource.target_project_ids == ["p_local"]
assert cfg.remote_datasource.target_project_ids == ["p_remote"]
with pytest.raises(AttributeError):
_ = cfg.target_project_ids
def test_datasource_target_ids_fallback_to_top_level_when_missing() -> None:
cfg = _build_config(top_level=["p_top"], local_ids=[], remote_ids=[])
def test_api_datasource_allows_empty_target_project_ids() -> None:
cfg = ApiDataSourceConfig(
type="api",
api_base_url="https://example.com",
)
local_ids, remote_ids = _resolve_datasource_target_project_ids(cfg)
assert local_ids == ["p_top"]
assert remote_ids == ["p_top"]
def test_datasource_target_ids_support_mixed_override_and_fallback() -> None:
cfg = _build_config(top_level=["p_top"], local_ids=["p_local"], remote_ids=[])
local_ids, remote_ids = _resolve_datasource_target_project_ids(cfg)
assert local_ids == ["p_local"]
assert remote_ids == ["p_top"]
def test_api_datasource_requires_target_project_ids() -> None:
with pytest.raises(ValidationError) as exc_info:
ApiDataSourceConfig(
type="api",
api_base_url="https://example.com",
target_project_ids=[], # Empty list should fail
)
assert "List should have at least 1 item" in str(exc_info.value)
with pytest.raises(ValidationError) as exc_info:
ApiDataSourceConfig(
type="api",
api_base_url="https://example.com",
# Missing target_project_ids should fail
)
assert "Field required" in str(exc_info.value) or "missing" in str(exc_info.value).lower()
assert cfg.target_project_ids == []
def test_api_datasource_rate_limit_config_fields() -> None:
@@ -102,11 +70,3 @@ def test_api_datasource_rate_limit_config_fields() -> None:
)
assert overridden.api_rate_limit_max_requests == 5
assert overridden.api_rate_limit_window_seconds == 2.5
def test_db_datasource_accepts_minimal_config() -> None:
cfg = DbDataSourceConfig(type="db")
assert cfg.type == "db"
assert cfg.target_project_ids == []
assert cfg.handler_configs == {}
+17 -6
View File
@@ -9,14 +9,15 @@ 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,
get_datasource_factory,
register_datasource_type,
)
from sync_state_machine.datasource import BaseDataSource
from sync_state_machine.datasource import ensure_builtin_datasource_types_registered
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
@@ -118,6 +119,13 @@ def _memory_factory(config: MemoryConfig, endpoint_name: str) -> MemoryDataSourc
return MemoryDataSource(namespace=f"{endpoint_name}:{config.namespace}")
def test_builtin_datasource_factories_are_registered() -> None:
ensure_builtin_datasource_types_registered(replace=True)
assert get_datasource_factory("jsonl") is not None
assert get_datasource_factory("api") is not None
@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)
@@ -126,7 +134,7 @@ async def test_create_pipeline_from_config_supports_registered_custom_datasource
config = PipelineRunConfig(
node_types=["memory_demo"],
local_datasource=MemoryConfig(type="memory", namespace="alpha"),
remote_datasource=DbDataSourceConfig(type="db"),
remote_datasource=MemoryConfig(type="memory", namespace="omega"),
strategies=[],
persist=PersistConfig(db_path=str(tmp_path / "memory-registry.db")),
logging=LoggingConfig(initialize=False),
@@ -136,8 +144,10 @@ async def test_create_pipeline_from_config_supports_registered_custom_datasource
try:
assert isinstance(pipeline.local_datasource, MemoryDataSource)
assert pipeline.local_datasource.namespace == "local:alpha"
assert isinstance(pipeline.remote_datasource, MemoryDataSource)
assert pipeline.remote_datasource.namespace == "remote:omega"
assert isinstance(pipeline.local_datasource.get_handler("memory_demo"), MemoryHandler)
assert isinstance(pipeline.remote_datasource.get_handler("memory_demo"), MemoryRemoteDbHandler)
assert isinstance(pipeline.remote_datasource.get_handler("memory_demo"), MemoryHandler)
finally:
await pipeline.close()
@@ -176,14 +186,15 @@ def test_build_config_supports_same_file_datasource_bootstrap(tmp_path: Path, mo
profile_path = tmp_path / "custom-memory.yaml"
profile_path.write_text(
"extensions:\n"
" datasource_bootstraps:\n"
" 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"
" type: jsonl\n"
" jsonl_dir: ${PROJECT_ROOT}/remote_datasource/datasource\n"
"persist:\n"
" db_path: ${PROJECT_ROOT}/runtime/test-memory-bootstrap.db\n"
"logging:\n"
@@ -195,4 +206,4 @@ def test_build_config_supports_same_file_datasource_bootstrap(tmp_path: Path, mo
assert config.local_datasource.type == "memory_bootstrap"
assert getattr(config.local_datasource, "namespace") == "local-demo"
assert config.remote_datasource.type == "db"
assert config.remote_datasource.type == "jsonl"
@@ -6,13 +6,25 @@ from pydantic import BaseModel
from sync_state_machine.common.registry import DomainRegistry
from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.config import DbDataSourceConfig, JsonlDataSourceConfig, LoggingConfig, PersistConfig, PipelineRunConfig
from sync_state_machine.config import (
BaseDataSourceConfig,
JsonlDataSourceConfig,
LoggingConfig,
PersistConfig,
PipelineRunConfig,
register_datasource_type,
)
from sync_state_machine.datasource.db import BaseDbHandler, DbDataSource
from sync_state_machine.datasource.jsonl import BaseJsonlHandler, JsonlDataSource
from sync_state_machine.pipeline.factory import _register_handlers
from sync_state_machine.sync_system.strategy import DefaultSyncStrategy
class TestRegisteredDbConfig(BaseDataSourceConfig):
__test__ = False
type: str = "test_db"
class TestDbSchema(BaseModel):
__test__ = False
id: str
@@ -56,7 +68,18 @@ class TestJsonlHandler(BaseJsonlHandler):
super().__init__(datasource, "test_db_handler", TestDbNode, TestDbSchema)
def _register_test_db_datasource_type() -> None:
register_datasource_type(
"test_db",
config_model=TestRegisteredDbConfig,
factory=lambda config, endpoint_name: DbDataSource(),
replace=True,
)
def test_register_handlers_uses_db_handler_for_db_datasource(tmp_path: Path) -> None:
_register_test_db_datasource_type()
if not DomainRegistry.is_registered("test_db_handler"):
DomainRegistry.register(
node_type="test_db_handler",
@@ -65,11 +88,11 @@ def test_register_handlers_uses_db_handler_for_db_datasource(tmp_path: Path) ->
strategy_class=TestDbStrategy,
jsonl_handler_class=TestJsonlHandler,
)
DomainRegistry.register_db_handler("test_db_handler", TestDbHandler)
DomainRegistry.register_handler("test_db_handler", "test_db", TestDbHandler)
config = PipelineRunConfig(
node_types=["test_db_handler"],
local_datasource=DbDataSourceConfig(type="db"),
local_datasource=TestRegisteredDbConfig(type="test_db"),
remote_datasource=JsonlDataSourceConfig(
type="jsonl",
jsonl_dir=str(tmp_path),
@@ -7,12 +7,23 @@ from pydantic import BaseModel
from sync_state_machine.common.registry import DomainRegistry
from sync_state_machine.common.sync_node import SyncNode
from sync_state_machine.config import DbDataSourceConfig, LoggingConfig, PersistConfig, PipelineRunConfig
from sync_state_machine.config import (
BaseDataSourceConfig,
LoggingConfig,
PersistConfig,
PipelineRunConfig,
register_datasource_type,
)
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 EndpointInitDbConfig(BaseDataSourceConfig):
__test__ = False
type: str = "endpoint_init_db"
class EndpointInitSchema(BaseModel):
__test__ = False
id: str
@@ -57,6 +68,7 @@ def _ensure_registered() -> None:
reg.node_class = EndpointInitNode
reg.strategy_class = EndpointInitStrategy
reg.db_handler_class = EndpointInitDbHandler
reg.handler_classes["endpoint_init_db"] = EndpointInitDbHandler
return
DomainRegistry.register(
@@ -66,6 +78,16 @@ def _ensure_registered() -> None:
strategy_class=EndpointInitStrategy,
db_handler_class=EndpointInitDbHandler,
)
DomainRegistry.register_handler(node_type, "endpoint_init_db", EndpointInitDbHandler)
def _register_endpoint_init_datasource_type() -> None:
register_datasource_type(
"endpoint_init_db",
config_model=EndpointInitDbConfig,
factory=lambda config, endpoint_name: DbDataSource(),
replace=True,
)
@pytest.mark.asyncio
@@ -74,16 +96,17 @@ async def test_factory_initializes_same_datasource_type_with_distinct_handler_co
monkeypatch: pytest.MonkeyPatch,
) -> None:
_ensure_registered()
_register_endpoint_init_datasource_type()
config = PipelineRunConfig(
node_types=["factory_endpoint_demo"],
local_datasource=DbDataSourceConfig(
type="db",
local_datasource=EndpointInitDbConfig(
type="endpoint_init_db",
handler_configs={"factory_endpoint_demo": {"side": "local", "batch_size": 10}},
target_project_ids=["project_local"],
),
remote_datasource=DbDataSourceConfig(
type="db",
remote_datasource=EndpointInitDbConfig(
type="endpoint_init_db",
handler_configs={"factory_endpoint_demo": {"side": "remote", "batch_size": 20}},
target_project_ids=["project_remote"],
),
@@ -4,6 +4,7 @@ from pathlib import Path
import pytest
from sync_state_machine.config import build_config_from_file
from sync_state_machine.config.run_presets import load_overrides_from_file
@@ -29,3 +30,97 @@ local_datasource:
assert overrides["logging"]["suppress_node_logs"] is True
assert str(tmp_path) in overrides["local_datasource"]["jsonl_dir"]
def test_load_overrides_executes_generic_extensions_bootstraps(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
bootstrap_module = tmp_path / "profile_bootstrap_demo.py"
bootstrap_module.write_text(
"BOOTSTRAP_CALLED = False\n"
"\n"
"def register():\n"
" global BOOTSTRAP_CALLED\n"
" BOOTSTRAP_CALLED = True\n",
encoding="utf-8",
)
monkeypatch.syspath_prepend(str(tmp_path))
yaml_file = tmp_path / "profile_with_bootstrap.yaml"
yaml_file.write_text(
"""
extensions:
bootstraps:
- profile_bootstrap_demo:register
local_datasource:
type: jsonl
jsonl_dir: ${PROJECT_ROOT}/filtered_datasource/datasource/filtered
remote_datasource:
type: jsonl
jsonl_dir: ${PROJECT_ROOT}/remote_datasource/datasource
persist:
db_path: ${PROJECT_ROOT}/runtime/test.db
logging:
initialize: false
""".strip(),
encoding="utf-8",
)
load_overrides_from_file(yaml_file, project_root=tmp_path)
import profile_bootstrap_demo
assert profile_bootstrap_demo.BOOTSTRAP_CALLED is True
def test_build_config_from_file_uses_registered_bootstraps(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
bootstrap_module = tmp_path / "config_from_file_bootstrap.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 FileBackedConfig(BaseDataSourceConfig):\n"
" model_config = ConfigDict(extra=\"forbid\", validate_assignment=True)\n"
" type: str = \"file_backed\"\n"
" namespace: str = \"demo\"\n"
"\n"
"class FileBackedDatasource(BaseDataSource):\n"
" pass\n"
"\n"
"def build_datasource(config, endpoint_name):\n"
" return FileBackedDatasource()\n"
"\n"
"def register():\n"
" register_datasource_type(\n"
" \"file_backed\",\n"
" config_model=FileBackedConfig,\n"
" factory=build_datasource,\n"
" replace=True,\n"
" )\n",
encoding="utf-8",
)
monkeypatch.syspath_prepend(str(tmp_path))
profile = tmp_path / "profile.yaml"
profile.write_text(
"""
extensions:
bootstraps:
- config_from_file_bootstrap:register
local_datasource:
type: file_backed
namespace: local-demo
remote_datasource:
type: jsonl
jsonl_dir: ${PROJECT_ROOT}/remote_datasource/datasource
persist:
db_path: ${PROJECT_ROOT}/runtime/test.db
logging:
initialize: false
""".strip(),
encoding="utf-8",
)
config = build_config_from_file(project_root=tmp_path, file_path=profile)
assert config.local_datasource.type == "file_backed"
assert getattr(config.local_datasource, "namespace") == "local-demo"