LangBot/tests/unit_tests/agent/conftest.py

"""Shared test fixtures for agent runner tests."""

from __future__ import annotations

import typing
from types import SimpleNamespace
from unittest.mock import AsyncMock

from langbot_plugin.entities.io.context import InstallationBinding


TEST_RUNTIME_BINDING = InstallationBinding(
    instance_uuid='instance-test',
    workspace_uuid='workspace-test',
    placement_generation=1,
    installation_uuid='00000000-0000-4000-8000-000000000001',
    runtime_revision=1,
    artifact_digest='a' * 64,
)


def bind_runtime_action_context(
    handler,
    application,
    *,
    plugin_identity: str = 'test/runner',
):
    """Simulate the trusted Runtime envelope used around direct action calls."""

    application.workspace_service = SimpleNamespace(
        get_execution_binding=AsyncMock(
            return_value=SimpleNamespace(
                instance_uuid=TEST_RUNTIME_BINDING.instance_uuid,
                workspace_uuid=TEST_RUNTIME_BINDING.workspace_uuid,
                placement_generation=TEST_RUNTIME_BINDING.placement_generation,
            )
        )
    )
    plugin_author, plugin_name = plugin_identity.split('/', 1)
    handler.register_installation_binding(
        TEST_RUNTIME_BINDING,
        plugin_author=plugin_author,
        plugin_name=plugin_name,
    )
    handler._current_action_context.set(TEST_RUNTIME_BINDING)
    return handler


def make_resources(
    models: list[dict] | None = None,
    tools: list[dict] | None = None,
    knowledge_bases: list[dict] | None = None,
    skills: list[dict] | None = None,
    storage: dict | None = None,
) -> dict[str, typing.Any]:
    """Create a minimal AgentResources dict for testing.

    Args:
        models: List of model dicts with 'model_id' key
        tools: List of tool dicts with 'tool_name' key
        knowledge_bases: List of KB dicts with 'kb_id' key
        skills: List of skill dicts with 'skill_name' key
        storage: Storage permissions dict
    Returns:
        AgentResources dict with all required fields
    """
    return {
        'models': models or [],
        'tools': tools or [],
        'knowledge_bases': knowledge_bases or [],
        'skills': skills or [],
        'storage': storage or {'plugin_storage': False, 'workspace_storage': False},
        'platform_capabilities': {},
    }


def make_session(
    run_id: str = 'test-run-id',
    runner_id: str = 'plugin:test/test-runner/default',
    query_id: int | None = 1,
    plugin_identity: str = 'test/test-runner',
    resources: dict | None = None,
    conversation_id: str | None = None,
    bot_id: str | None = None,
    workspace_id: str | None = None,
    thread_id: str | None = None,
    available_apis: dict[str, bool] | None = None,
    state_policy: dict[str, typing.Any] | None = None,
    state_context: dict[str, typing.Any] | None = None,
    execution_query: typing.Any | None = None,
) -> dict[str, typing.Any]:
    """Create a minimal AgentRunSession dict for testing.

    Args:
        run_id: Unique run identifier
        runner_id: Runner descriptor ID
        query_id: Host entry query ID
        plugin_identity: Plugin identifier (author/name)
        resources: AgentResources dict (uses make_resources() default if None)

    Returns:
        AgentRunSession dict with run-scoped authorization snapshot
    """
    import time

    now = int(time.time())
    res = resources if resources is not None else make_resources()
    apis = available_apis if available_apis is not None else {}
    policy = (
        state_policy if state_policy is not None else {'enable_state': True, 'state_scopes': ['conversation', 'actor']}
    )
    context = state_context if state_context is not None else {}

    authorized_ids: dict[str, set[str]] = {
        'model': {m.get('model_id') for m in res.get('models', [])},
        'tool': {t.get('tool_name') for t in res.get('tools', [])},
        'knowledge_base': {kb.get('kb_id') for kb in res.get('knowledge_bases', [])},
        'skill': {s.get('skill_name') for s in res.get('skills', [])},
    }
    authorized_operations: dict[str, dict[str, set[str]]] = {
        'model': {
            m.get('model_id'): set(m.get('operations') or ['invoke', 'stream', 'rerank', 'count_tokens'])
            for m in res.get('models', [])
            if m.get('model_id')
        },
        'tool': {
            t.get('tool_name'): set(t.get('operations') or ['detail', 'call'])
            for t in res.get('tools', [])
            if t.get('tool_name')
        },
        'knowledge_base': {
            kb.get('kb_id'): set(kb.get('operations') or ['list', 'retrieve'])
            for kb in res.get('knowledge_bases', [])
            if kb.get('kb_id')
        },
        'skill': {
            s.get('skill_name'): set(s.get('operations') or ['activate'])
            for s in res.get('skills', [])
            if s.get('skill_name')
        },
    }

    return {
        'run_id': run_id,
        'runner_id': runner_id,
        'query_id': query_id,
        'execution_query': execution_query,
        'plugin_identity': plugin_identity,
        'authorization': {
            'resources': res,
            'available_apis': apis,
            'conversation_id': conversation_id,
            'bot_id': bot_id,
            'workspace_id': workspace_id,
            'thread_id': thread_id,
            'state_policy': policy,
            'state_context': context,
            'authorized_ids': authorized_ids,
            'authorized_operations': authorized_operations,
        },
        'status': {
            'started_at': now,
            'last_activity_at': now,
        },
    }