Files
Auto-GPT/tests/integration/test_execute_code.py
Reinier van der Leer 3fe2246468 Agent loop v2: Prompting improvements & WIP planning (#5077)
* Add categories to command registry

* Fix tests

* Clean up prompt generation

* Rename Performance Evaluations to Best Practices
* Move specification of response format from system prompt to Agent.construct_base_prompt
* Clean up PromptGenerator class

* Add debug logging to AIConfig autogeneration

* Clarify prompting and add support for multiple thought processes to Agent

* WIP: PlanningAgent

* Disable message history by default on BaseAgent

* Add CommandOutput and ThoughtProcessOutput type aliases

* Fix interrupts in main.py

* Use custom exceptions and clean up exception/error handling

* Remove duplicate agent_history.py

* Update PlanningAgent from upstream

* WIP: Support for dynamic in-prompt context

* WIP: response formats for PlanningAgent three-stage cycle

* Remove browsing overlay & separate browsing from extraction code

* Fix human feedback

* Fix tests

* Include history in Agent prompt generation

* Code improvements in agent.py

* Add ask_user command and revise system prompt
2023-08-19 17:44:50 +02:00

136 lines
4.5 KiB
Python

import os
import random
import string
import tempfile
import pytest
import autogpt.commands.execute_code as sut # system under testing
from autogpt.agents.agent import Agent
from autogpt.agents.utils.exceptions import (
AccessDeniedError,
InvalidArgumentError,
OperationNotAllowedError,
)
from autogpt.config import Config
@pytest.fixture
def random_code(random_string) -> str:
return f"print('Hello {random_string}!')"
@pytest.fixture
def python_test_file(config: Config, random_code: str) -> str:
temp_file = tempfile.NamedTemporaryFile(dir=config.workspace_path, suffix=".py")
temp_file.write(str.encode(random_code))
temp_file.flush()
yield temp_file.name
temp_file.close()
@pytest.fixture
def random_string():
return "".join(random.choice(string.ascii_lowercase) for _ in range(10))
def test_execute_python_file(python_test_file: str, random_string: str, agent: Agent):
result: str = sut.execute_python_file(python_test_file, agent=agent)
assert result.replace("\r", "") == f"Hello {random_string}!\n"
def test_execute_python_code(random_code: str, random_string: str, agent: Agent):
ai_name = agent.ai_config.ai_name
result: str = sut.execute_python_code(random_code, "test_code", agent=agent)
assert result.replace("\r", "") == f"Hello {random_string}!\n"
# Check that the code is stored
destination = os.path.join(
agent.config.workspace_path, ai_name, "executed_code", "test_code.py"
)
with open(destination) as f:
assert f.read() == random_code
def test_execute_python_code_disallows_name_arg_path_traversal(
random_code: str, agent: Agent
):
with pytest.raises(AccessDeniedError, match="path traversal"):
sut.execute_python_code(random_code, name="../../test_code", agent=agent)
# Check that the code is not stored in parent directory
dst_with_traversal = agent.workspace.get_path("test_code.py")
assert not dst_with_traversal.is_file(), "Path traversal by filename not prevented"
def test_execute_python_code_overwrites_file(random_code: str, agent: Agent):
ai_name = agent.ai_config.ai_name
destination = os.path.join(
agent.config.workspace_path, ai_name, "executed_code", "test_code.py"
)
os.makedirs(os.path.dirname(destination), exist_ok=True)
with open(destination, "w+") as f:
f.write("This will be overwritten")
sut.execute_python_code(random_code, "test_code.py", agent=agent)
# Check that the file is updated with the new code
with open(destination) as f:
assert f.read() == random_code
def test_execute_python_file_invalid(agent: Agent):
with pytest.raises(InvalidArgumentError):
sut.execute_python_file("not_python", agent)
def test_execute_python_file_not_found(agent: Agent):
with pytest.raises(
FileNotFoundError,
match=r"python: can't open file '([a-zA-Z]:)?[/\\\-\w]*notexist.py': \[Errno 2\] No such file or directory",
):
sut.execute_python_file("notexist.py", agent)
def test_execute_shell(random_string: str, agent: Agent):
result = sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
assert f"Hello {random_string}!" in result
def test_execute_shell_local_commands_not_allowed(random_string: str, agent: Agent):
result = sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
assert f"Hello {random_string}!" in result
def test_execute_shell_denylist_should_deny(agent: Agent, random_string: str):
agent.config.shell_denylist = ["echo"]
with pytest.raises(OperationNotAllowedError, match="not allowed"):
sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
def test_execute_shell_denylist_should_allow(agent: Agent, random_string: str):
agent.config.shell_denylist = ["cat"]
result = sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
assert "Hello" in result and random_string in result
def test_execute_shell_allowlist_should_deny(agent: Agent, random_string: str):
agent.config.shell_command_control = sut.ALLOWLIST_CONTROL
agent.config.shell_allowlist = ["cat"]
with pytest.raises(OperationNotAllowedError, match="not allowed"):
sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
def test_execute_shell_allowlist_should_allow(agent: Agent, random_string: str):
agent.config.shell_command_control = sut.ALLOWLIST_CONTROL
agent.config.shell_allowlist = ["echo"]
result = sut.execute_shell(f"echo 'Hello {random_string}!'", agent)
assert "Hello" in result and random_string in result