""" Unit tests for executor tenant_id/agent_id injection. Tests that execute_tool injects tenant_id and agent_id into handler kwargs before calling the handler, so context-aware tools (kb_search, calendar_lookup) receive tenant context without the LLM needing to provide it. """ from __future__ import annotations import uuid from typing import Any from unittest.mock import AsyncMock, MagicMock import pytest def _make_tool(handler: Any, requires_confirmation: bool = False) -> Any: """Create a minimal ToolDefinition-like object for tests.""" tool = MagicMock() tool.handler = handler tool.requires_confirmation = requires_confirmation tool.parameters = { "type": "object", "properties": { "query": {"type": "string"}, }, "required": ["query"], } return tool class TestExecutorTenantInjection: @pytest.mark.asyncio async def test_tenant_id_injected_into_handler_kwargs(self) -> None: """Handler should receive tenant_id even though LLM didn't provide it.""" from orchestrator.tools.executor import execute_tool received_kwargs: dict[str, Any] = {} async def mock_handler(**kwargs: Any) -> str: received_kwargs.update(kwargs) return "handler result" tool = _make_tool(mock_handler) registry = {"test_tool": tool} tenant_id = uuid.uuid4() agent_id = uuid.uuid4() audit_logger = MagicMock() audit_logger.log_tool_call = AsyncMock() tool_call = { "function": { "name": "test_tool", "arguments": '{"query": "hello world"}', } } result = await execute_tool(tool_call, registry, tenant_id, agent_id, audit_logger) assert result == "handler result" assert "tenant_id" in received_kwargs assert received_kwargs["tenant_id"] == str(tenant_id) @pytest.mark.asyncio async def test_agent_id_injected_into_handler_kwargs(self) -> None: """Handler should receive agent_id even though LLM didn't provide it.""" from orchestrator.tools.executor import execute_tool received_kwargs: dict[str, Any] = {} async def mock_handler(**kwargs: Any) -> str: received_kwargs.update(kwargs) return "ok" tool = _make_tool(mock_handler) registry = {"test_tool": tool} tenant_id = uuid.uuid4() agent_id = uuid.uuid4() audit_logger = MagicMock() audit_logger.log_tool_call = AsyncMock() tool_call = { "function": { "name": "test_tool", "arguments": '{"query": "test"}', } } await execute_tool(tool_call, registry, tenant_id, agent_id, audit_logger) assert "agent_id" in received_kwargs assert received_kwargs["agent_id"] == str(agent_id) @pytest.mark.asyncio async def test_injected_ids_are_strings(self) -> None: """Injected tenant_id and agent_id should be strings, not UUIDs.""" from orchestrator.tools.executor import execute_tool received_kwargs: dict[str, Any] = {} async def mock_handler(**kwargs: Any) -> str: received_kwargs.update(kwargs) return "ok" tool = _make_tool(mock_handler) registry = {"test_tool": tool} tenant_id = uuid.uuid4() agent_id = uuid.uuid4() audit_logger = MagicMock() audit_logger.log_tool_call = AsyncMock() tool_call = { "function": { "name": "test_tool", "arguments": '{"query": "test"}', } } await execute_tool(tool_call, registry, tenant_id, agent_id, audit_logger) assert isinstance(received_kwargs["tenant_id"], str) assert isinstance(received_kwargs["agent_id"], str) @pytest.mark.asyncio async def test_llm_provided_args_preserved(self) -> None: """Original LLM-provided args should still be present after injection.""" from orchestrator.tools.executor import execute_tool received_kwargs: dict[str, Any] = {} async def mock_handler(**kwargs: Any) -> str: received_kwargs.update(kwargs) return "ok" tool = _make_tool(mock_handler) registry = {"test_tool": tool} tenant_id = uuid.uuid4() agent_id = uuid.uuid4() audit_logger = MagicMock() audit_logger.log_tool_call = AsyncMock() tool_call = { "function": { "name": "test_tool", "arguments": '{"query": "search term from LLM"}', } } await execute_tool(tool_call, registry, tenant_id, agent_id, audit_logger) assert received_kwargs["query"] == "search term from LLM" assert received_kwargs["tenant_id"] == str(tenant_id) assert received_kwargs["agent_id"] == str(agent_id) @pytest.mark.asyncio async def test_injection_after_schema_validation(self) -> None: """Injection happens after validation — injected keys don't cause schema failures.""" from orchestrator.tools.executor import execute_tool # Tool requires exactly 'query', nothing else in schema required # Schema should pass even though we inject tenant_id/agent_id async def mock_handler(**kwargs: Any) -> str: return "passed" tool = _make_tool(mock_handler) registry = {"test_tool": tool} tenant_id = uuid.uuid4() agent_id = uuid.uuid4() audit_logger = MagicMock() audit_logger.log_tool_call = AsyncMock() tool_call = { "function": { "name": "test_tool", "arguments": '{"query": "test"}', } } result = await execute_tool(tool_call, registry, tenant_id, agent_id, audit_logger) assert result == "passed"