strepsiades
127 lines
3.9 KiB
Python
127 lines
3.9 KiB
Python
from __future__ import annotations
|
|
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
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 (
|
|
BaseDataSourceConfig,
|
|
LoggingConfig,
|
|
PersistConfig,
|
|
PipelineRunConfig,
|
|
register_datasource_type,
|
|
)
|
|
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
|
|
|
|
|
|
class EndpointInitConfig(BaseDataSourceConfig):
|
|
__test__ = False
|
|
type: str = "endpoint_init_memory"
|
|
|
|
|
|
class EndpointInitSchema(BaseModel):
|
|
__test__ = False
|
|
id: str
|
|
name: str | None = None
|
|
|
|
|
|
class EndpointInitNode(SyncNode[EndpointInitSchema]):
|
|
__test__ = False
|
|
node_type = "factory_endpoint_demo"
|
|
schema = EndpointInitSchema
|
|
|
|
|
|
class EndpointInitStrategy(DefaultSyncStrategy[EndpointInitSchema]):
|
|
__test__ = False
|
|
schema = EndpointInitSchema
|
|
|
|
|
|
class EndpointInitHandler(BaseNodeHandler[EndpointInitSchema]):
|
|
__test__ = False
|
|
|
|
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", ""))
|
|
|
|
async def load(self):
|
|
return []
|
|
|
|
async def create_all(self, nodes):
|
|
return []
|
|
|
|
async def update_all(self, nodes):
|
|
return []
|
|
|
|
|
|
def _ensure_registered() -> None:
|
|
node_type = "factory_endpoint_demo"
|
|
if DomainRegistry.is_registered(node_type):
|
|
reg = DomainRegistry.get_registration(node_type)
|
|
reg.schema = EndpointInitSchema
|
|
reg.node_class = EndpointInitNode
|
|
reg.strategy_class = EndpointInitStrategy
|
|
reg.handler_classes["endpoint_init_memory"] = EndpointInitHandler
|
|
return
|
|
|
|
DomainRegistry.register(
|
|
node_type=node_type,
|
|
schema=EndpointInitSchema,
|
|
node_class=EndpointInitNode,
|
|
strategy_class=EndpointInitStrategy,
|
|
)
|
|
DomainRegistry.register_handler(node_type, "endpoint_init_memory", EndpointInitHandler)
|
|
|
|
|
|
def _register_endpoint_init_datasource_type() -> None:
|
|
register_datasource_type(
|
|
"endpoint_init_memory",
|
|
config_model=EndpointInitConfig,
|
|
factory=lambda config, endpoint_name: BaseDataSource(),
|
|
replace=True,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_factory_initializes_same_datasource_type_with_distinct_handler_configs(
|
|
tmp_path: Path,
|
|
monkeypatch: pytest.MonkeyPatch,
|
|
) -> None:
|
|
_ensure_registered()
|
|
_register_endpoint_init_datasource_type()
|
|
|
|
config = PipelineRunConfig(
|
|
node_types=["factory_endpoint_demo"],
|
|
local_datasource=EndpointInitConfig(
|
|
type="endpoint_init_memory",
|
|
handler_configs={"factory_endpoint_demo": {"side": "local", "batch_size": 10}},
|
|
),
|
|
remote_datasource=EndpointInitConfig(
|
|
type="endpoint_init_memory",
|
|
handler_configs={"factory_endpoint_demo": {"side": "remote", "batch_size": 20}},
|
|
),
|
|
strategies=[],
|
|
persist=PersistConfig(db_path=str(tmp_path / "factory-init.db")),
|
|
logging=LoggingConfig(initialize=False),
|
|
)
|
|
|
|
pipeline = await create_pipeline_from_config(config)
|
|
try:
|
|
local_handler = pipeline.local_datasource.get_handler("factory_endpoint_demo")
|
|
remote_handler = pipeline.remote_datasource.get_handler("factory_endpoint_demo")
|
|
|
|
assert isinstance(local_handler, EndpointInitHandler)
|
|
assert isinstance(remote_handler, EndpointInitHandler)
|
|
assert local_handler.handler_config == {"side": "local", "batch_size": 10}
|
|
assert remote_handler.handler_config == {"side": "remote", "batch_size": 20}
|
|
finally:
|
|
await pipeline.close()
|