zk-data-agent/tests/test_python_exec_tool.py

from __future__ import annotations

import sys
import tempfile
from pathlib import Path
from unittest import TestCase

from src.agent_tools import build_tool_context, default_tool_registry, execute_tool
from src.agent_types import AgentPermissions, AgentRuntimeConfig


class PythonExecToolTests(TestCase):
    def test_python_exec_runs_inline_code_with_shell_permission(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            config = AgentRuntimeConfig(
                cwd=Path(tmp_dir),
                permissions=AgentPermissions(allow_shell_commands=True),
            )
            context = build_tool_context(config)
            result = execute_tool(
                default_tool_registry(),
                'python_exec',
                {'code': 'print("hello from python")'},
                context,
            )

        self.assertTrue(result.ok)
        self.assertIn('exit_code=0', result.content)
        self.assertIn('hello from python', result.content)
        self.assertEqual(result.metadata.get('action'), 'python_exec')
        self.assertEqual(result.metadata.get('mode'), 'code')

    def test_python_exec_exposes_session_scratchpad(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            scratchpad = (root / 'session' / 'scratchpad').resolve()
            scratchpad.mkdir(parents=True)
            config = AgentRuntimeConfig(
                cwd=root,
                permissions=AgentPermissions(allow_shell_commands=True),
            )
            context = build_tool_context(config, scratchpad_directory=scratchpad)
            result = execute_tool(
                default_tool_registry(),
                'python_exec',
                {'code': 'import os\nprint(os.environ["PYTHON_EXEC_SCRATCHPAD"])'},
                context,
            )

        self.assertTrue(result.ok)
        self.assertIn(str(scratchpad), result.content)
        self.assertEqual(result.metadata.get('scratchpad_directory'), str(scratchpad))

    def test_python_exec_prefers_user_python_env(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            user_env = (root / 'account' / 'python' / '.venv').resolve()
            bin_dir = user_env / 'bin'
            bin_dir.mkdir(parents=True)
            python_bin = bin_dir / 'python'
            try:
                python_bin.symlink_to(sys.executable)
            except OSError:
                python_bin.write_text(
                    '#!/bin/sh\nexec "$PYTHON_EXEC_TEST_REAL_PYTHON" "$@"\n',
                    encoding='utf-8',
                )
                python_bin.chmod(0o755)
            config = AgentRuntimeConfig(
                cwd=root,
                permissions=AgentPermissions(allow_shell_commands=True),
                python_env_dir=user_env,
            )
            context = build_tool_context(
                config,
                extra_env={'PYTHON_EXEC_TEST_REAL_PYTHON': sys.executable},
            )
            result = execute_tool(
                default_tool_registry(),
                'python_exec',
                {'code': 'import os, sys\nprint(sys.executable)\nprint(os.environ["PYTHON_EXEC_VENV"])'},
                context,
            )

        self.assertTrue(result.ok)
        self.assertIn(str(user_env), result.content)
        self.assertEqual(result.metadata.get('python_env_dir'), str(user_env))

    def test_python_package_show_uses_user_python_env(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            user_env = (root / 'account' / 'python' / '.venv').resolve()
            bin_dir = user_env / 'bin'
            bin_dir.mkdir(parents=True)
            python_bin = bin_dir / 'python'
            try:
                python_bin.symlink_to(sys.executable)
            except OSError:
                python_bin.write_text(
                    '#!/bin/sh\nexec "$PYTHON_EXEC_TEST_REAL_PYTHON" "$@"\n',
                    encoding='utf-8',
                )
                python_bin.chmod(0o755)
            config = AgentRuntimeConfig(
                cwd=root,
                permissions=AgentPermissions(allow_shell_commands=True),
                python_env_dir=user_env,
            )
            context = build_tool_context(
                config,
                extra_env={'PYTHON_EXEC_TEST_REAL_PYTHON': sys.executable},
            )
            result = execute_tool(
                default_tool_registry(),
                'python_package',
                {'action': 'show'},
                context,
            )

        self.assertTrue(result.ok)
        self.assertIn(str(user_env), result.content)
        self.assertEqual(result.metadata.get('python_env_dir'), str(user_env))

    def test_python_exec_is_blocked_without_shell_permission(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            config = AgentRuntimeConfig(cwd=Path(tmp_dir))
            context = build_tool_context(config)
            result = execute_tool(
                default_tool_registry(),
                'python_exec',
                {'code': 'print("blocked")'},
                context,
            )

        self.assertFalse(result.ok)
        self.assertIn('--allow-shell', result.content)
        self.assertEqual(result.metadata.get('error_kind'), 'permission_denied')

    def test_python_exec_runs_workspace_script_with_args(self) -> None:
        with tempfile.TemporaryDirectory() as tmp_dir:
            root = Path(tmp_dir)
            script = root / 'tool_script.py'
            script.write_text(
                'import sys\nprint("|".join(sys.argv[1:]))\n',
                encoding='utf-8',
            )
            config = AgentRuntimeConfig(
                cwd=root,
                permissions=AgentPermissions(allow_shell_commands=True),
            )
            context = build_tool_context(config)
            result = execute_tool(
                default_tool_registry(),
                'python_exec',
                {'script_path': 'tool_script.py', 'args': ['a', 'b']},
                context,
            )

        self.assertTrue(result.ok)
        self.assertIn('a|b', result.content)
        self.assertEqual(result.metadata.get('mode'), 'script')