"""Tests for MCP-specific security rules."""

import ast

from mcpaudit.rules.mcp_tool_safety import MCPToolDescriptionRule, MCPEnvAccessRule


class TestMCPToolDescription:
    def setup_method(self):
        self.rule = MCPToolDescriptionRule()

    def _check(self, code: str):
        tree = ast.parse(code)
        return self.rule.check(tree, code)

    def test_tool_with_docstring_passes(self):
        code = '''
def handle_tool_read(path):
    """Reads a file from disk."""
    return open(path).read()
'''
        assert len(self._check(code)) == 0

    def test_tool_without_docstring_detected(self):
        code = '''
def handle_tool_read(path):
    return open(path).read()
'''
        findings = self._check(code)
        assert len(findings) == 1
        assert findings[0].rule_id == "MCP-001"

    def test_non_tool_function_ignored(self):
        code = '''
def helper_function(x):
    return x * 2
'''
        assert len(self._check(code)) == 0

    def test_decorated_tool_without_docstring(self):
        code = '''
@server.tool
def my_func(path):
    return path
'''
        findings = self._check(code)
        assert len(findings) == 1

    def test_async_tool_without_docstring(self):
        code = '''
async def handle_tool_search(query):
    return []
'''
        findings = self._check(code)
        assert len(findings) == 1

    def test_multiple_tools_mixed(self):
        code = '''
def handle_tool_a():
    """Has docstring."""
    pass

def handle_tool_b():
    pass
'''
        findings = self._check(code)
        assert len(findings) == 1
        assert "handle_tool_b" in findings[0].description

    def test_rule_metadata(self):
        assert self.rule.rule_id == "MCP-001"
        assert self.rule.name == "MCP Tool Missing Description"


class TestMCPEnvAccess:
    def setup_method(self):
        self.rule = MCPEnvAccessRule()

    def _check(self, code: str):
        tree = ast.parse(code)
        return self.rule.check(tree, code)

    def test_env_access_in_tool_detected(self):
        code = '''
import os
def handle_tool_config():
    return os.getenv("API_KEY")
'''
        findings = self._check(code)
        assert len(findings) == 1
        assert findings[0].rule_id == "MCP-002"

    def test_env_access_outside_tool_ignored(self):
        code = '''
import os
API_KEY = os.getenv("API_KEY")
def process(x):
    return x
'''
        assert len(self._check(code)) == 0

    def test_os_environ_in_tool_detected(self):
        code = '''
import os
def call_tool(name, args):
    key = os.environ("SECRET")
    return key
'''
        findings = self._check(code)
        assert len(findings) == 1

    def test_no_env_access_passes(self):
        code = '''
def handle_tool_echo(text):
    return text
'''
        assert len(self._check(code)) == 0

    def test_rule_metadata(self):
        assert self.rule.rule_id == "MCP-002"
        assert self.rule.name == "Direct Environment Access in Tool"
