mirror of
https://github.com/langbot-app/LangBot.git
synced 2026-08-09 12:40:59 +00:00
huanghuoguoguo
80 lines
2.7 KiB
Python
80 lines
2.7 KiB
Python
from __future__ import annotations
|
|
|
|
from types import SimpleNamespace
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
from langbot.pkg.api.http.context import ExecutionContext
|
|
from langbot.pkg.agent.runner.errors import RunnerNotFoundError
|
|
from langbot.pkg.pipeline.controller import Controller
|
|
|
|
|
|
def make_app():
|
|
app = SimpleNamespace()
|
|
app.instance_config = SimpleNamespace(data={'concurrency': {'pipeline': 10}})
|
|
app.logger = MagicMock()
|
|
app.pipeline_mgr = SimpleNamespace()
|
|
app.pipeline_mgr.get_pipeline_by_uuid = AsyncMock()
|
|
app.sess_mgr = SimpleNamespace()
|
|
app.sess_mgr.get_session = AsyncMock(return_value=SimpleNamespace())
|
|
app.agent_run_orchestrator = SimpleNamespace()
|
|
app.agent_run_orchestrator.try_claim_steering_from_query = AsyncMock()
|
|
return app
|
|
|
|
|
|
def make_pipeline():
|
|
return SimpleNamespace(
|
|
pipeline_entity=SimpleNamespace(config={'ai': {'runner': {'id': 'plugin:test/runner/default'}}}),
|
|
bound_plugins=['test/runner'],
|
|
bound_mcp_servers=[],
|
|
)
|
|
|
|
|
|
def make_query(query_id: int, pipeline_uuid: str):
|
|
context = ExecutionContext(
|
|
instance_uuid='instance-test',
|
|
workspace_uuid='workspace-test',
|
|
placement_generation=1,
|
|
pipeline_uuid=pipeline_uuid,
|
|
)
|
|
return SimpleNamespace(
|
|
query_id=query_id,
|
|
pipeline_uuid=pipeline_uuid,
|
|
variables={},
|
|
_execution_context=context,
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_try_claim_steering_returns_false_when_runner_lookup_fails():
|
|
app = make_app()
|
|
app.pipeline_mgr.get_pipeline_by_uuid.return_value = make_pipeline()
|
|
app.agent_run_orchestrator.try_claim_steering_from_query.side_effect = RunnerNotFoundError(
|
|
'plugin:missing/runner/default'
|
|
)
|
|
controller = Controller(app)
|
|
query = make_query(1, 'pipeline-001')
|
|
|
|
claimed = await controller._try_claim_steering_before_session_slot(query)
|
|
|
|
assert claimed is False
|
|
app.logger.warning.assert_called_once()
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_try_claim_steering_sets_pipeline_context_before_claiming():
|
|
app = make_app()
|
|
pipeline = make_pipeline()
|
|
app.pipeline_mgr.get_pipeline_by_uuid.return_value = pipeline
|
|
app.agent_run_orchestrator.try_claim_steering_from_query.return_value = True
|
|
controller = Controller(app)
|
|
query = make_query(2, 'pipeline-002')
|
|
|
|
claimed = await controller._try_claim_steering_before_session_slot(query)
|
|
|
|
assert claimed is True
|
|
assert query.pipeline_config is pipeline.pipeline_entity.config
|
|
assert query.variables['_pipeline_bound_plugins'] == ['test/runner']
|
|
app.agent_run_orchestrator.try_claim_steering_from_query.assert_awaited_once_with(query)
|