ss-tools/backend/tests/api/test_assistant_tool_registry.py

# #region Test.Assistant.ToolRegistry [C:2] [TYPE Module] [SEMANTICS test,assistant,registry,tool]
# @BRIEF Tests for _tool_registry.py — core decorator, catalog, dispatch, permissions.
# @RELATION BINDS_TO -> [AssistantToolRegistry]
# @TEST_CONTRACT: assistant_tool decorator -> registers tool in _tools dict, sets __assistant_tool__ attr
# @TEST_CONTRACT: AssistantTool dataclass -> stores all metadata attributes
# @TEST_CONTRACT: get_safe_ops -> returns risk_level=='safe' operations + dataset_review_answer_context
# @TEST_CONTRACT: get_permission_checks -> returns only tools with non-empty permission_checks
# @TEST_CONTRACT: get_catalog -> filters by user permissions, includes defaults
# @TEST_CONTRACT: dispatch -> routes to correct handler, raises 400 for unknown ops
# @TEST_CONTRACT: _check_any_permission -> passes on first valid check, raises 403 on all fail
# @TEST_CONTRACT: _has_any_permission -> boolean wrapper around _check_any_permission
# @TEST_EDGE: unknown_operation -> HTTPException 400
# @TEST_EDGE: no_permission -> HTTPException 403
# @TEST_EDGE: catalog_filtered_by_permissions -> excluded tools not in catalog
# @TEST_EDGE: dataset_review_dispatch_path -> special-case routing
# @TEST_EDGE: get_safe_ops_includes_registry_and_extra -> both sources merged

from datetime import UTC, datetime
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from fastapi import HTTPException

from src.schemas.auth import User


@pytest.fixture
def current_user():
    return User(
        id="user-1", username="testuser", email="test@example.com",
        is_active=True, auth_source="internal",
        created_at=datetime.now(UTC), last_login=datetime.now(UTC), roles=[],
    )


@pytest.fixture
def config_manager():
    cm = MagicMock()
    cm.get_environments = MagicMock(return_value=[])
    return cm


@pytest.fixture
def task_manager():
    tm = MagicMock()
    tm.create_task = AsyncMock(return_value=MagicMock(id="task-42"))
    tm.get_task = MagicMock(return_value=None)
    tm.get_tasks = MagicMock(return_value=[])
    return tm


@pytest.fixture
def db():
    return MagicMock()


@pytest.fixture(autouse=True)
def clear_registry():
    """Clear _tools registry before and after each test."""
    import src.api.routes.assistant._tool_registry as reg
    saved = dict(reg._tools)
    reg._tools.clear()
    yield
    reg._tools.clear()
    reg._tools.update(saved)


class TestAssistantToolDataclass:
    """AssistantTool — metadata container."""

    def test_minimal(self):
        from src.api.routes.assistant._tool_registry import AssistantTool

        async def dummy_handler(**kwargs):
            return ("ok", None, [])

        tool = AssistantTool(
            operation="test_op",
            domain="test",
            description="A test tool",
            handler=dummy_handler,
        )
        assert tool.operation == "test_op"
        assert tool.risk_level == "safe"
        assert tool.requires_confirmation is False
        assert tool.permission_checks == []
        assert tool.defaults is None

    def test_full(self):
        from src.api.routes.assistant._tool_registry import AssistantTool

        async def dummy(**kwargs):
            return ("ok", None, [])

        tool = AssistantTool(
            operation="risky_op",
            domain="admin",
            description="Risky",
            handler=dummy,
            required_entities=["env"],
            optional_entities=["dashboard_id"],
            risk_level="guarded",
            requires_confirmation=True,
            permission_checks=[("admin", "EXECUTE")],
            defaults={"timeout": 30},
        )
        assert tool.risk_level == "guarded"
        assert tool.requires_confirmation is True
        assert tool.permission_checks == [("admin", "EXECUTE")]
        assert tool.defaults == {"timeout": 30}


class TestAssistantToolDecorator:
    """assistant_tool decorator — registration mechanics."""

    @pytest.mark.asyncio
    async def test_registers_tool(self):
        from src.api.routes.assistant._tool_registry import assistant_tool, _tools

        @assistant_tool(operation="my_op", domain="test", description="My tool")
        async def my_handler(**kwargs):
            return ("done", None, [])

        assert "my_op" in _tools
        assert _tools["my_op"].operation == "my_op"
        assert _tools["my_op"].domain == "test"

    @pytest.mark.asyncio
    async def test_sets_handler_attr(self):
        from src.api.routes.assistant._tool_registry import assistant_tool

        @assistant_tool(operation="op_with_attr", domain="test", description="Check attr")
        async def handler_with_attr(**kwargs):
            return ("done", None, [])

        assert hasattr(handler_with_attr, "__assistant_tool__")
        assert handler_with_attr.__assistant_tool__.operation == "op_with_attr"

    @pytest.mark.asyncio
    async def test_registers_with_defaults(self):
        from src.api.routes.assistant._tool_registry import assistant_tool, _tools

        @assistant_tool(
            operation="with_defaults",
            domain="test",
            description="Has defaults",
            required_entities=["x"],
            optional_entities=["y"],
            risk_level="guarded",
            requires_confirmation=True,
            permission_checks=[("res", "action")],
            defaults={"key": "val"},
        )
        async def handler_defaults(**kwargs):
            return ("done", None, [])

        tool = _tools["with_defaults"]
        assert tool.required_entities == ["x"]
        assert tool.risk_level == "guarded"
        assert tool.requires_confirmation is True
        assert tool.permission_checks == [("res", "action")]
        assert tool.defaults == {"key": "val"}


class TestGetSafeOps:
    """get_safe_ops — safe operations set."""

    def test_empty_registry_returns_minimal(self):
        from src.api.routes.assistant._tool_registry import get_safe_ops

        ops = get_safe_ops()
        assert "dataset_review_answer_context" in ops
        # No registered tools yet — only the hardcoded extra
        assert len(ops) >= 1

    def test_includes_safe_tools(self):
        from src.api.routes.assistant._tool_registry import assistant_tool, get_safe_ops

        @assistant_tool(operation="safe_op", domain="t", description="s", risk_level="safe")
        async def safe_h(**kwargs):
            return ("ok", None, [])

        @assistant_tool(operation="guarded_op", domain="t", description="g", risk_level="guarded")
        async def guarded_h(**kwargs):
            return ("ok", None, [])

        ops = get_safe_ops()
        assert "safe_op" in ops
        assert "guarded_op" not in ops
        assert "dataset_review_answer_context" in ops


class TestGetPermissionChecks:
    """get_permission_checks — permission mapping."""

    def test_no_tools_with_checks(self):
        from src.api.routes.assistant._tool_registry import get_permission_checks

        checks = get_permission_checks()
        # Should include dataset_review ops from _schemas
        assert isinstance(checks, dict)

    def test_includes_tool_checks(self):
        from src.api.routes.assistant._tool_registry import assistant_tool, get_permission_checks

        @assistant_tool(
            operation="protected_op", domain="t", description="p",
            permission_checks=[("admin", "WRITE")],
        )
        async def protected_h(**kwargs):
            return ("ok", None, [])

        @assistant_tool(operation="public_op", domain="t", description="safe")
        async def public_h(**kwargs):
            return ("ok", None, [])

        checks = get_permission_checks()
        assert "protected_op" in checks
        assert checks["protected_op"] == [("admin", "WRITE")]
        # public_op has no checks, should not appear
        assert "public_op" not in checks


class TestCatalog:
    """get_catalog — LLM tool catalog generation."""

    def test_empty_registry(self, current_user, config_manager, db):
        from src.api.routes.assistant._tool_registry import get_catalog

        catalog = get_catalog(current_user, config_manager, db)
        assert catalog == []

    def test_excludes_protected_tools_without_permission(self, current_user, config_manager, db):
        from src.api.routes.assistant._tool_registry import assistant_tool, get_catalog, _has_any_permission

        @assistant_tool(
            operation="admin_only", domain="admin", description="Admin tool",
            permission_checks=[("admin", "WRITE")],
        )
        async def admin_h(**kwargs):
            return ("ok", None, [])

        @assistant_tool(
            operation="public", domain="public", description="Public tool",
        )
        async def public_h(**kwargs):
            return ("ok", None, [])

        # Mock _has_any_permission to return False for admin_only
        with patch("src.api.routes.assistant._tool_registry._has_any_permission",
                   side_effect=lambda u, checks: checks != [("admin", "WRITE")]):
            catalog = get_catalog(current_user, config_manager, db)

        ops = [c["operation"] for c in catalog]
        assert "public" in ops
        assert "admin_only" not in ops

    def test_includes_defaults_in_entry(self, current_user, config_manager, db):
        from src.api.routes.assistant._tool_registry import assistant_tool, get_catalog

        @assistant_tool(
            operation="with_defaults", domain="t", description="d",
            defaults={"timeout": 60},
        )
        async def h(**kwargs):
            return ("ok", None, [])

        catalog = get_catalog(current_user, config_manager, db)
        assert len(catalog) == 1
        assert catalog[0]["defaults"] == {"timeout": 60}

    def test_entries_sorted(self, current_user, config_manager, db):
        from src.api.routes.assistant._tool_registry import assistant_tool, get_catalog

        @assistant_tool(operation="z_tool", domain="t", description="z")
        async def z_h(**kwargs):
            return ("ok", None, [])

        @assistant_tool(operation="a_tool", domain="t", description="a")
        async def a_h(**kwargs):
            return ("ok", None, [])

        catalog = get_catalog(current_user, config_manager, db)
        ops = [c["operation"] for c in catalog]
        assert ops == sorted(ops)


class TestCheckAnyPermission:
    """_check_any_permission and _has_any_permission."""

    def test_any_permission_passes(self, current_user):
        from src.api.routes.assistant._tool_registry import _check_any_permission

        with patch("src.api.routes.assistant._tool_registry.has_permission") as mock_hp:
            mock_hp.return_value = lambda u: None  # passes
            _check_any_permission(current_user, [("res", "READ")])
            mock_hp.assert_called_once_with("res", "READ")

    def test_any_permission_first_wins(self, current_user):
        from src.api.routes.assistant._tool_registry import _check_any_permission

        call_order = []

        def first_check(u):
            call_order.append("first")
            raise HTTPException(status_code=403, detail="denied")

        def second_check(u):
            call_order.append("second")
            return None

        with patch("src.api.routes.assistant._tool_registry.has_permission") as mock_hp:
            mock_hp.side_effect = [first_check, second_check]
            _check_any_permission(current_user, [("res1", "READ"), ("res2", "WRITE")])
            assert call_order == ["first", "second"]

    def test_all_permissions_fail_raises_403(self, current_user):
        from src.api.routes.assistant._tool_registry import _check_any_permission

        with patch("src.api.routes.assistant._tool_registry.has_permission") as mock_hp:
            mock_hp.return_value = lambda u: (_ for _ in ()).throw(
                HTTPException(status_code=403, detail="denied")
            )
            with pytest.raises(HTTPException) as exc:
                _check_any_permission(current_user, [("res1", "X"), ("res2", "Y")])
            assert exc.value.status_code == 403

    def test_empty_checks_raises_403(self, current_user):
        from src.api.routes.assistant._tool_registry import _check_any_permission

        with pytest.raises(HTTPException) as exc:
            _check_any_permission(current_user, [])
        assert exc.value.status_code == 403

    def test_has_any_permission_true(self, current_user):
        from src.api.routes.assistant._tool_registry import _has_any_permission

        with patch("src.api.routes.assistant._tool_registry._check_any_permission",
                   return_value=None):
            assert _has_any_permission(current_user, [("res", "READ")]) is True

    def test_has_any_permission_false(self, current_user):
        from src.api.routes.assistant._tool_registry import _has_any_permission

        with patch("src.api.routes.assistant._tool_registry._check_any_permission",
                   side_effect=HTTPException(403)):
            assert _has_any_permission(current_user, [("res", "READ")]) is False


class TestDispatch:
    """dispatch — operation routing."""

    @pytest.mark.asyncio
    async def test_dispatches_to_registered_handler(
        self, current_user, task_manager, config_manager, db
    ):
        from src.api.routes.assistant._tool_registry import assistant_tool, dispatch, _tools

        @assistant_tool(operation="my_dispatch_op", domain="test", description="Dispatch test")
        async def my_handler(intent, current_user, task_manager, config_manager, db):
            return ("executed", None, [])

        text, task_id, actions = await dispatch(
            "my_dispatch_op", {"entities": {}}, current_user,
            task_manager, config_manager, db,
        )
        assert text == "executed"

    @pytest.mark.asyncio
    async def test_unknown_operation_raises_400(
        self, current_user, task_manager, config_manager, db
    ):
        from src.api.routes.assistant._tool_registry import dispatch

        with pytest.raises(HTTPException) as exc:
            await dispatch(
                "nonexistent_op", {"entities": {}}, current_user,
                task_manager, config_manager, db,
            )
        assert exc.value.status_code == 400
        assert "Unsupported operation" in exc.value.detail

    @pytest.mark.asyncio
    async def test_dataset_review_ops_go_to_subdispatcher(
        self, current_user, task_manager, config_manager, db
    ):
        from src.api.routes.assistant._tool_registry import dispatch

        # The sub-dispatcher is imported inside dispatch() via lazy import
        # Patch at the source module level — the lazy import will pick it up
        with patch(
            "src.api.routes.assistant._dataset_review_dispatch._dispatch_dataset_review_intent",
            new=AsyncMock(return_value=("review done", None, [])),
        ) as mock_sub:
            text, task_id, actions = await dispatch(
                "dataset_review_approve_mappings",
                {"entities": {"dataset_review_session_id": "s1", "session_version": 1}},
                current_user, task_manager, config_manager, db,
            )
            assert text == "review done"
            mock_sub.assert_awaited_once()

    @pytest.mark.asyncio
    async def test_dataset_review_answer_context_goes_to_subdispatcher(
        self, current_user, task_manager, config_manager, db
    ):
        from src.api.routes.assistant._tool_registry import dispatch

        with patch(
            "src.api.routes.assistant._dataset_review_dispatch._dispatch_dataset_review_intent",
            new=AsyncMock(return_value=("context summary", None, [])),
        ) as mock_sub:
            text, task_id, actions = await dispatch(
                "dataset_review_answer_context",
                {"entities": {"dataset_review_session_id": "s1", "session_version": 1}},
                current_user, task_manager, config_manager, db,
            )
            assert text == "context summary"
            mock_sub.assert_awaited_once()


# #endregion Test.Assistant.ToolRegistry