From 7e3af1842af244a54cdc8acdba57b75c4796b958 Mon Sep 17 00:00:00 2001 From: Mery-Sanz Date: Thu, 10 Apr 2025 09:57:35 +0200 Subject: [PATCH] add coments and docstring --- tests/agents/test_agent_one_tool.py | 29 +++++++++--- ...est_agent_prompt_system_master_template.py | 24 +++++++--- .../tools/test_tool_generic_linux_command.py | 45 ++++++++++++++++--- 3 files changed, 81 insertions(+), 17 deletions(-) diff --git a/tests/agents/test_agent_one_tool.py b/tests/agents/test_agent_one_tool.py index fac7fda9..2bb3c1f0 100644 --- a/tests/agents/test_agent_one_tool.py +++ b/tests/agents/test_agent_one_tool.py @@ -1,3 +1,10 @@ +""" +This module contains tests for the one-tool agent functionality, specifically +for the CTF agent. It includes tests to verify the agent's instructions and +configuration, as well as its ability to execute a Linux command using the +generic_linux_command tool. +""" + import pytest from tests.fake_model import FakeModel from tests.test_responses import ( @@ -11,20 +18,26 @@ from cai.agents.one_tool import one_tool_agent @pytest.mark.asyncio async def test_ctf_agent_instructions_and_configuration(): + """Test the CTF agent's instructions and configuration.""" agent = transfer_to_one_tool_agent() - # Check that the agent has the generic_linux_command tool - assert any(tool.name== "generic_linux_command" for tool in agent.tools) - - # Optionally, you can check the agent's instructions and configuration + # Check if the agent has the expected tool + assert any(tool.name == "generic_linux_command" for tool in agent.tools) + + # Ensure the agent has instructions set assert agent.instructions is not None + + # Verify the agent's name assert agent.name == "CTF agent" @pytest.mark.asyncio async def test_ctf_agent_executes_linux_command(): + """Test the CTF agent's ability to execute a Linux command.""" model = FakeModel() agent = transfer_to_one_tool_agent() agent.model = model + + # Set up the model's expected outputs for the command execution model.add_multiple_turn_outputs( [ [ @@ -37,10 +50,14 @@ async def test_ctf_agent_executes_linux_command(): ] ) + # Run the agent with a specific input result = await Runner.run(agent, input="List files") + # Verify the final output of the command execution assert result.final_output == "result of the command: flag{12345}" + + # Ensure the number of raw responses is as expected assert len(result.raw_responses) == 2 - assert any("generic_linux_command" in str(item) for item in result.to_input_list()) - + # Check if the command tool was used in the input list + assert any("generic_linux_command" in str(item) for item in result.to_input_list()) \ No newline at end of file diff --git a/tests/agents/test_agent_prompt_system_master_template.py b/tests/agents/test_agent_prompt_system_master_template.py index cb345760..ef440e86 100644 --- a/tests/agents/test_agent_prompt_system_master_template.py +++ b/tests/agents/test_agent_prompt_system_master_template.py @@ -1,33 +1,47 @@ +""" +This module contains tests for the Mako template rendering of the system master template +used in the agent framework. It includes tests to verify the correct rendering of the +template with various configurations, including the presence of agent instructions and +handling of environment variables. +""" + import os import pytest from mako.template import Template +# Fixture to load the Mako template for the system master template @pytest.fixture def template(): return Template(filename="src/cai/prompts/core/system_master_template.md") +# Fixture to create a base agent with predefined instructions @pytest.fixture def base_agent(): return type('Agent', (), {'instructions': 'Test instructions'})() def test_master_template_basic(template, base_agent): - """Test basic master template rendering without optional components""" + """Test basic master template rendering without optional components.""" result = template.render(agent=base_agent, reasoning_content=None, ctf_instructions="") print(result) + # Verify that the agent's instructions are included in the rendered template assert 'Test instructions' in result + # Ensure that the CTF_INSIDE placeholder is not present assert 'CTF_INSIDE' not in result def test_master_template_with_env_vars(template, base_agent): - """Test master template with environment variables and vector DB""" + """Test master template with environment variables and vector DB.""" + # Set an environment variable for the CTF name os.environ['CTF_NAME'] = 'test_ctf' result = template.render(agent=base_agent, reasoning_content=None, ctf_instructions="") - print(result) + # Verify that the agent's instructions are included in the rendered template assert "Test instructions" in result + # Clean up by deleting the environment variable del os.environ['CTF_NAME'] def test_master_template_no_instructions(template): - """Test master template without agent instructions""" + """Test master template without agent instructions.""" + # Create an agent with empty instructions agent = type('Agent', (), {'instructions': ''})() result = template.render(agent=agent, reasoning_content=None, ctf_instructions="") - print(result) + # Verify that the rendered template starts with an empty string assert result.strip().startswith('') diff --git a/tests/tools/test_tool_generic_linux_command.py b/tests/tools/test_tool_generic_linux_command.py index 8b6b3209..97c152fb 100644 --- a/tests/tools/test_tool_generic_linux_command.py +++ b/tests/tools/test_tool_generic_linux_command.py @@ -1,18 +1,51 @@ +""" +This module contains tests for the generic Linux command tool functionality. +It includes tests for executing regular commands and handling invalid commands. +""" + import pytest import json -import asyncio from unittest.mock import MagicMock -from unittest.mock import patch from cai.tools.reconnaissance.generic_linux_command import generic_linux_command - async def test_generic_linux_command_regular_commands(): - mock_ctx = MagicMock() + """Test the execution of a regular command using the generic Linux command tool.""" + mock_ctx = MagicMock() # Create a mock context for the command execution params = { - "command": "echo", - "args": "'hello'" + "command": "echo", # Command to be executed + "args": "'hello'" # Arguments for the command } + # Invoke the tool with the specified parameters and await the result result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params)) + # Assert that the result matches the expected output assert result.replace("\n", "") == 'hello' + +async def test_generic_linux_command_ls(): + """Test the execution of the 'ls' command using the generic Linux command tool.""" + mock_ctx = MagicMock() # Create a mock context for the command execution + params = { + "command": "ls", # Command to be executed + "args": "-l" # Arguments for the command + } + + # Invoke the tool with the specified parameters and await the result + result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params)) + + # Assert that the output contains 'total', which is typical for 'ls -l' + assert "total" in result + +async def test_generic_linux_command_invalid_command(): + """Test the handling of an invalid command using the generic Linux command tool.""" + mock_ctx = MagicMock() # Create a mock context for the command execution + params = { + "command": "invalid_command", # Invalid command to be executed + "args": "" # No arguments for the command + } + + # Invoke the tool with the specified parameters and await the result + result = await generic_linux_command.on_invoke_tool(mock_ctx, json.dumps(params)) + + # Assert that the result indicates the command was not found + assert "command not found" in result