greydgl--pentestgpt/tests/integration/test_controller.py

"""Tests for AgentController.

Integration tests for the agent controller lifecycle management.
"""

from pathlib import Path

import pytest

from pentestgpt.core.backend import AgentMessage, MessageType
from pentestgpt.core.config import PentestGPTConfig
from pentestgpt.core.controller import AgentController, AgentState
from pentestgpt.core.session import SessionStore
from tests.conftest import MockBackend


@pytest.mark.integration
class TestAgentController:
    """Tests for AgentController."""

    @pytest.fixture
    def config(self, temp_working_dir: Path) -> PentestGPTConfig:
        """Create test config."""
        return PentestGPTConfig(
            target="test.example.com",
            working_directory=temp_working_dir,
        )

    @pytest.fixture
    def controller(
        self,
        config: PentestGPTConfig,
        mock_backend: MockBackend,
        temp_sessions_dir: Path,
    ) -> AgentController:
        """Create controller with mock backend."""
        session_store = SessionStore(sessions_dir=temp_sessions_dir)
        return AgentController(
            config=config,
            backend=mock_backend,
            session_store=session_store,
        )

    def test_initial_state(self, controller: AgentController):
        """Test controller starts in idle state."""
        assert controller.state == AgentState.IDLE

    @pytest.mark.asyncio
    async def test_run_success(self, controller: AgentController, mock_backend: MockBackend):
        """Test successful run."""
        # Set up messages
        mock_backend.set_messages(
            [
                AgentMessage(type=MessageType.TEXT, content="Starting task..."),
                AgentMessage(type=MessageType.TEXT, content="Found flag{test123}!"),
                AgentMessage(type=MessageType.RESULT, content=None, metadata={"cost_usd": 0.5}),
            ]
        )

        result = await controller.run("Test task")

        assert result["success"] is True
        assert "flag{test123}" in result["flags_found"]
        assert result["session_id"] is not None

    @pytest.mark.asyncio
    async def test_run_with_tool_messages(
        self, controller: AgentController, mock_backend: MockBackend
    ):
        """Test run with tool messages."""
        mock_backend.set_messages(
            [
                AgentMessage(
                    type=MessageType.TOOL_START,
                    content=None,
                    tool_name="bash",
                    tool_args={"command": "nmap -sV target"},
                ),
                AgentMessage(
                    type=MessageType.TOOL_RESULT,
                    content="PORT STATE SERVICE\n22 open ssh",
                    tool_name="bash",
                ),
                AgentMessage(type=MessageType.RESULT, content=None, metadata={}),
            ]
        )

        result = await controller.run("Scan target")
        assert result["success"] is True

    @pytest.mark.asyncio
    @pytest.mark.slow
    @pytest.mark.skip(reason="Async pause/resume test requires concurrent task management")
    async def test_pause_and_resume(self, controller: AgentController, mock_backend: MockBackend):
        """Test pause and resume functionality.

        Note: This test is skipped because properly testing async pause/resume
        requires running the controller in a separate task and sending control
        signals concurrently. The current mock approach cannot reliably test this.
        """
        pass

    def test_pause_request(self, controller: AgentController):
        """Test pause request in wrong state."""
        # Can't pause when idle
        assert controller.pause() is False

    def test_resume_request(self, controller: AgentController):
        """Test resume request in wrong state."""
        # Can't resume when idle
        assert controller.resume() is False

    def test_stop_request(self, controller: AgentController):
        """Test stop request."""
        assert controller.stop() is True

    @pytest.mark.asyncio
    @pytest.mark.slow
    @pytest.mark.skip(reason="Async stop test requires concurrent task management")
    async def test_stop_during_run(self, controller: AgentController, mock_backend: MockBackend):
        """Test stopping during run.

        Note: This test is skipped because properly testing async stop
        requires running the controller in a separate task and sending stop
        signals concurrently. The current mock approach cannot reliably test this.
        """
        pass

    def test_inject_instruction(self, controller: AgentController):
        """Test inject instruction."""
        # Can't inject when idle
        assert controller.inject("New instruction") is False

    @pytest.mark.asyncio
    async def test_flag_detection(self, controller: AgentController, mock_backend: MockBackend):
        """Test flag detection patterns."""
        mock_backend.set_messages(
            [
                AgentMessage(type=MessageType.TEXT, content="Found flag{abc123}"),
                AgentMessage(type=MessageType.TEXT, content="Also HTB{test_flag}"),
                AgentMessage(
                    type=MessageType.TEXT,
                    content="And 0123456789abcdef0123456789abcdef",
                ),
                AgentMessage(type=MessageType.RESULT, content=None, metadata={}),
            ]
        )

        result = await controller.run("Find flags")

        assert len(result["flags_found"]) == 3
        assert "flag{abc123}" in result["flags_found"]
        assert "HTB{test_flag}" in result["flags_found"]

    @pytest.mark.asyncio
    async def test_session_not_found(self, controller: AgentController):
        """Test resuming non-existent session."""
        result = await controller.run("Task", resume_session_id="nonexistent")
        assert result["success"] is False
        assert "not found" in result["error"]

    @pytest.mark.asyncio
    async def test_backend_error(self, controller: AgentController, mock_backend: MockBackend):
        """Test handling backend errors."""

        async def error_messages():
            raise Exception("Backend error")
            yield  # Make it a generator

        mock_backend.receive_messages = error_messages  # type: ignore[method-assign]

        result = await controller.run("Test task")
        assert result["success"] is False
        assert "Backend error" in result["error"]