Back to Repositories

Testing SimpleAgent Code Generation and Improvement in gpt-engineer

This test suite validates the SimpleAgent class implementation in gpt-engineer, focusing on initialization and improvement capabilities for code generation and modification. The tests verify the agent’s ability to create and modify Python programs based on natural language prompts.

Test Coverage Overview

The test suite provides comprehensive coverage of SimpleAgent’s core functionality:

Tests initialization process with file creation and content verification
Validates code improvement capabilities through string manipulation
Covers file system operations and execution environment interactions
Tests integration with mock AI responses and prompt handling

Implementation Analysis

The testing approach utilizes pytest’s framework features for unit testing the SimpleAgent class. It employs temporary directories for file operations and implements mock AI responses to simulate the agent’s behavior. The tests follow a pattern of setup, execution, and assertion to validate both the initial code generation and subsequent modifications.

Technical Details

Key technical components include:

pytest for test framework
tempfile for temporary directory management
MockAI class for AI response simulation
DiskExecutionEnv for file system operations
FilesDict for managing code artifacts

Best Practices Demonstrated

The test suite demonstrates several testing best practices:

Isolation of tests using temporary directories
Mocking of external dependencies
Clear test case organization
Verification of both positive paths and modifications
Proper cleanup of test resources

gpt-engineer-org/gpt-engineer

tests/core/default/test_simple_agent.py

            
import tempfile

import pytest

from langchain.schema import AIMessage

from gpt_engineer.core.default.disk_execution_env import DiskExecutionEnv
from gpt_engineer.core.default.paths import ENTRYPOINT_FILE
from gpt_engineer.core.default.simple_agent import SimpleAgent
from gpt_engineer.core.files_dict import FilesDict
from gpt_engineer.core.prompt import Prompt
from tests.mock_ai import MockAI


def test_init():
    temp_dir = tempfile.mkdtemp()
    mock_ai = MockAI(
        [
            AIMessage(
                "hello_world.py
```
with open('output.txt', 'w') as file:
    file.write('Hello World!')
```"
            ),
            AIMessage("```run.sh
python3 hello_world.py
```"),
        ],
    )
    lean_agent = SimpleAgent.with_default_config(temp_dir, mock_ai)
    outfile = "output.txt"
    code = lean_agent.init(
        Prompt(
            f"Make a program that prints 'Hello World!' to a file called '{outfile}'"
        )
    )

    env = DiskExecutionEnv()
    env.upload(code).run(f"bash {ENTRYPOINT_FILE}")
    code = env.download()

    assert outfile in code
    assert code[outfile] == "Hello World!"


def test_improve():
    temp_dir = tempfile.mkdtemp()
    code = FilesDict(
        {
            "main.py": "def write_hello_world_to_file(filename):
    \"\"\"
    Writes 'Hello World!' to the specified file.
    
    :param filename: The name of the file to write to.
    \"\"\"
    with open(filename, 'w') as file:
        file.write('Hello World!')

if __name__ == \"__main__\":
    output_filename = 'output.txt'
    write_hello_world_to_file(output_filename)",
            "requirements.txt": "# No dependencies required",
            "run.sh": "python3 main.py
",
        }
    )
    mock_ai = MockAI(
        [
            AIMessage(
                "```diff
--- main.py
+++ main.py
@@ -7,3 +7,3 @@
     with open(filename, 'w') as file:
-        file.write('Hello World!')
+        file.write('!dlroW olleH')
```"
            )
        ]
    )
    lean_agent = SimpleAgent.with_default_config(temp_dir, mock_ai)
    code = lean_agent.improve(
        code,
        Prompt(
            "Change the program so that it prints '!dlroW olleH' instead of 'Hello World!' "
        ),
        f"bash {ENTRYPOINT_FILE}",
    )

    env = DiskExecutionEnv()
    env.upload(code).run(f"bash {ENTRYPOINT_FILE}")
    code = env.download()

    outfile = "output.txt"
    assert outfile in code
    assert code[outfile] == "!dlroW olleH"


if __name__ == "__main__":
    pytest.main()