Make test write file hard (#4481)

tests/challenges/basic_abilities/test_write_file.py

@@ -1,5 +1,6 @@
+from typing import List
+
 import pytest
-from pytest_mock import MockerFixture
 
 from autogpt.agent import Agent
 from autogpt.commands.file_operations import read_file
@@ -8,21 +9,34 @@ from tests.challenges.challenge_decorator.challenge_decorator import challenge
 from tests.challenges.utils import run_interaction_loop
 from tests.utils import requires_api_key
 
-CYCLE_COUNT = 3
+CYCLE_COUNT_PER_LEVEL = [1, 1]
+EXPECTED_OUTPUTS_PER_LEVEL = [
+    {"hello_world.txt": ["Hello World"]},
+    {"hello_world_1.txt": ["Hello World"], "hello_world_2.txt": ["Hello World"]},
+]
 
 
 @requires_api_key("OPENAI_API_KEY")
 @pytest.mark.vcr
 @challenge
 def test_write_file(
-    writer_agent: Agent,
-    patched_api_requestor: MockerFixture,
+    file_system_agents: List[Agent],
+    patched_api_requestor: None,
     monkeypatch: pytest.MonkeyPatch,
     config: Config,
     level_to_run: int,
 ) -> None:
-    file_path = str(writer_agent.workspace.get_path("hello_world.txt"))
-    run_interaction_loop(monkeypatch, writer_agent, CYCLE_COUNT)
+    file_system_agent = file_system_agents[level_to_run - 1]
+    run_interaction_loop(
+        monkeypatch, file_system_agent, CYCLE_COUNT_PER_LEVEL[level_to_run - 1]
+    )
 
-    content = read_file(file_path, config)
-    assert content == "Hello World", f"Expected 'Hello World', got {content}"
+    expected_outputs = EXPECTED_OUTPUTS_PER_LEVEL[level_to_run - 1]
+
+    for file_name, expected_lines in expected_outputs.items():
+        file_path = str(file_system_agent.workspace.get_path(file_name))
+        content = read_file(file_path, config)
+        for expected_line in expected_lines:
+            assert (
+                expected_line in content
+            ), f"Expected '{expected_line}' in file {file_name}, but it was not found"

tests/challenges/current_score.json

@@ -5,7 +5,7 @@
             "max_level_beaten": 1
         },
         "write_file": {
-            "max_level": 1,
+            "max_level": 2,
             "max_level_beaten": 1
         }
     },

tests/integration/agent_factory.py

@@ -70,41 +70,39 @@ def browser_agent(agent_test_config, memory_none: NoMemory, workspace: Workspace
 
 
 @pytest.fixture
-def writer_agent(agent_test_config, memory_none: NoMemory, workspace: Workspace):
-    command_registry = CommandRegistry()
-    command_registry.import_commands("autogpt.commands.file_operations")
-    command_registry.import_commands("autogpt.app")
-    command_registry.import_commands("autogpt.commands.task_statuses")
+def file_system_agents(
+    agent_test_config, memory_json_file: NoMemory, workspace: Workspace
+):
+    agents = []
+    command_registry = get_command_registry(agent_test_config)
 
-    ai_config = AIConfig(
-        ai_name="write_to_file-GPT",
-        ai_role="an AI designed to use the write_to_file command to write 'Hello World' into a file named \"hello_world.txt\" and then use the task_complete command to complete the task.",
-        ai_goals=[
-            "Use the write_to_file command to write 'Hello World' into a file named \"hello_world.txt\".",
-            "Use the task_complete command to complete the task.",
-            "Do not use any other commands.",
-        ],
-    )
-    ai_config.command_registry = command_registry
+    ai_goals = [
+        "Write 'Hello World' into a file named \"hello_world.txt\".",
+        'Write \'Hello World\' into 2 files named "hello_world_1.txt"and "hello_world_2.txt".',
+    ]
 
-    triggering_prompt = (
-        "Determine which next command to use, and respond using the"
-        " format specified above:"
-    )
-    system_prompt = ai_config.construct_full_prompt()
-
-    agent = Agent(
-        ai_name="",
-        memory=memory_none,
-        command_registry=command_registry,
-        config=ai_config,
-        next_action_count=0,
-        system_prompt=system_prompt,
-        triggering_prompt=triggering_prompt,
-        workspace_directory=workspace.root,
-    )
-
-    return agent
+    for ai_goal in ai_goals:
+        ai_config = AIConfig(
+            ai_name="File System Agent",
+            ai_role="an AI designed to manage a file system.",
+            ai_goals=[ai_goal],
+        )
+        ai_config.command_registry = command_registry
+        system_prompt = ai_config.construct_full_prompt()
+        Config().set_continuous_mode(False)
+        agents.append(
+            Agent(
+                ai_name="Information Retrieval Agent",
+                memory=memory_json_file,
+                command_registry=command_registry,
+                config=ai_config,
+                next_action_count=0,
+                system_prompt=system_prompt,
+                triggering_prompt=DEFAULT_TRIGGERING_PROMPT,
+                workspace_directory=workspace.root,
+            )
+        )
+    return agents
 
 
 @pytest.fixture
@@ -145,15 +143,8 @@ def information_retrieval_agents(
     agent_test_config, memory_json_file, workspace: Workspace
 ):
     agents = []
-    command_registry = CommandRegistry()
-    enabled_command_categories = [
-        x
-        for x in COMMAND_CATEGORIES
-        if x not in agent_test_config.disabled_command_categories
-    ]
+    command_registry = get_command_registry(agent_test_config)
 
-    for command_category in enabled_command_categories:
-        command_registry.import_commands(command_category)
     ai_goals = [
         "Write to a file called output.txt tesla's revenue in 2022 after searching for 'tesla revenue 2022'.",
         "Write to a file called output.txt tesla's revenue in 2022.",
@@ -284,3 +275,15 @@ def debug_code_agent(agent_test_config, memory_json_file, workspace: Workspace):
     )
 
     return agent
+
+
+def get_command_registry(agent_test_config):
+    command_registry = CommandRegistry()
+    enabled_command_categories = [
+        x
+        for x in COMMAND_CATEGORIES
+        if x not in agent_test_config.disabled_command_categories
+    ]
+    for command_category in enabled_command_categories:
+        command_registry.import_commands(command_category)
+    return command_registry