Add comprehensive unit tests for AWEL agent operators and context

Co-authored-by: Aries-ckt <13723926+Aries-ckt@users.noreply.github.com>
2026-07-17 01:58:47 +00:00 · 2026-01-30 15:20:08 +00:00
parent b40db652ed
commit 515a16bc59
3 changed files with 1323 additions and 0 deletions
--- a/tests/unit_tests/agent/init.py
+++ b/tests/unit_tests/agent/init.py
@@ -0,0 +1 @@
+"""Unit tests for the agent module."""
--- a/tests/unit_tests/agent/test_agent_context.py
+++ b/tests/unit_tests/agent/test_agent_context.py
@@ -0,0 +1,576 @@
+"""Unit tests for agent context classes.
+
+Tests for AgentContext, AgentGenerateContext, AgentMessage, and AgentReviewInfo
+from dbgpt.agent.core.agent module.
+"""
+
+from unittest.mock import MagicMock
+
+import pytest
+
+from dbgpt.agent.core.action.base import ActionOutput
+from dbgpt.agent.core.agent import (
+    Agent,
+    AgentContext,
+    AgentGenerateContext,
+    AgentMessage,
+    AgentReviewInfo,
+)
+from dbgpt.core.interface.message import ModelMessageRoleType
+
+
+class TestAgentContext:
+    """Test suite for AgentContext class."""
+
+    def test_initialization_with_defaults(self):
+        """Test AgentContext initialization with default values."""
+        context = AgentContext(conv_id="test_conv_123")
+
+        assert context.conv_id == "test_conv_123"
+        assert context.gpts_app_code is None
+        assert context.gpts_app_name is None
+        assert context.language is None
+        assert context.max_chat_round == 100
+        assert context.max_retry_round == 10
+        assert context.max_new_tokens == 1024
+        assert context.temperature == 0.5
+        assert context.allow_format_str_template is False
+        assert context.verbose is False
+        assert context.app_link_start is False
+        assert context.enable_vis_message is True
+
+    def test_initialization_with_custom_values(self):
+        """Test AgentContext initialization with custom values."""
+        context = AgentContext(
+            conv_id="custom_conv",
+            gpts_app_code="app_001",
+            gpts_app_name="TestApp",
+            language="zh",
+            max_chat_round=50,
+            max_retry_round=5,
+            max_new_tokens=2048,
+            temperature=0.7,
+            allow_format_str_template=True,
+            verbose=True,
+            app_link_start=True,
+            enable_vis_message=False,
+        )
+
+        assert context.conv_id == "custom_conv"
+        assert context.gpts_app_code == "app_001"
+        assert context.gpts_app_name == "TestApp"
+        assert context.language == "zh"
+        assert context.max_chat_round == 50
+        assert context.max_retry_round == 5
+        assert context.max_new_tokens == 2048
+        assert context.temperature == 0.7
+        assert context.allow_format_str_template is True
+        assert context.verbose is True
+        assert context.app_link_start is True
+        assert context.enable_vis_message is False
+
+    def test_to_dict(self):
+        """Test AgentContext.to_dict() method."""
+        context = AgentContext(
+            conv_id="test_conv",
+            gpts_app_code="app_001",
+            language="en",
+            max_chat_round=50,
+        )
+
+        result = context.to_dict()
+
+        assert isinstance(result, dict)
+        assert result["conv_id"] == "test_conv"
+        assert result["gpts_app_code"] == "app_001"
+        assert result["language"] == "en"
+        assert result["max_chat_round"] == 50
+        assert "max_retry_round" in result
+        assert "temperature" in result
+
+
+class TestAgentGenerateContext:
+    """Test suite for AgentGenerateContext class."""
+
+    def test_initialization_with_minimal_values(self):
+        """Test AgentGenerateContext with minimal required values."""
+        mock_agent = MagicMock(spec=Agent)
+        mock_message = AgentMessage(content="test message")
+
+        context = AgentGenerateContext(message=mock_message, sender=mock_agent)
+
+        assert context.message == mock_message
+        assert context.sender == mock_agent
+        assert context.reviewer is None
+        assert context.silent is False
+        assert context.already_failed is False
+        assert context.last_speaker is None
+        assert context.already_started is False
+        assert context.begin_agent is None
+        assert context.rely_messages == []
+        assert context.final is True
+        assert context.memory is None
+        assert context.agent_context is None
+        assert context.llm_client is None
+        assert context.round_index is None
+
+    def test_initialization_with_full_values(self):
+        """Test AgentGenerateContext with all values."""
+        mock_sender = MagicMock(spec=Agent)
+        mock_reviewer = MagicMock(spec=Agent)
+        mock_last_speaker = MagicMock(spec=Agent)
+        mock_message = AgentMessage(content="test message")
+        rely_msg1 = AgentMessage(content="rely message 1")
+        rely_msg2 = AgentMessage(content="rely message 2")
+        mock_agent_context = AgentContext(conv_id="test_conv")
+
+        context = AgentGenerateContext(
+            message=mock_message,
+            sender=mock_sender,
+            reviewer=mock_reviewer,
+            silent=True,
+            already_failed=True,
+            last_speaker=mock_last_speaker,
+            already_started=True,
+            begin_agent="begin_agent_name",
+            rely_messages=[rely_msg1, rely_msg2],
+            final=False,
+            agent_context=mock_agent_context,
+            round_index=5,
+        )
+
+        assert context.message == mock_message
+        assert context.sender == mock_sender
+        assert context.reviewer == mock_reviewer
+        assert context.silent is True
+        assert context.already_failed is True
+        assert context.last_speaker == mock_last_speaker
+        assert context.already_started is True
+        assert context.begin_agent == "begin_agent_name"
+        assert len(context.rely_messages) == 2
+        assert context.rely_messages[0] == rely_msg1
+        assert context.rely_messages[1] == rely_msg2
+        assert context.final is False
+        assert context.agent_context == mock_agent_context
+        assert context.round_index == 5
+
+    def test_rely_messages_handling(self):
+        """Test AgentGenerateContext handling of rely_messages."""
+        mock_agent = MagicMock(spec=Agent)
+        mock_message = AgentMessage(content="main message")
+        rely_msg1 = AgentMessage(content="rely 1", role=ModelMessageRoleType.HUMAN)
+        rely_msg2 = AgentMessage(content="rely 2", role=ModelMessageRoleType.AI)
+
+        context = AgentGenerateContext(
+            message=mock_message,
+            sender=mock_agent,
+            rely_messages=[rely_msg1, rely_msg2],
+        )
+
+        assert len(context.rely_messages) == 2
+        assert context.rely_messages[0].content == "rely 1"
+        assert context.rely_messages[1].content == "rely 2"
+
+    def test_already_failed_flag(self):
+        """Test AgentGenerateContext already_failed flag."""
+        mock_agent = MagicMock(spec=Agent)
+        mock_message = AgentMessage(content="test")
+
+        # Test with already_failed=False (default)
+        context1 = AgentGenerateContext(message=mock_message, sender=mock_agent)
+        assert context1.already_failed is False
+
+        # Test with already_failed=True
+        context2 = AgentGenerateContext(
+            message=mock_message, sender=mock_agent, already_failed=True
+        )
+        assert context2.already_failed is True
+
+    def test_to_dict(self):
+        """Test AgentGenerateContext.to_dict() method."""
+        mock_agent = MagicMock(spec=Agent)
+        mock_message = AgentMessage(content="test message")
+        mock_agent_context = AgentContext(conv_id="test_conv")
+
+        context = AgentGenerateContext(
+            message=mock_message,
+            sender=mock_agent,
+            agent_context=mock_agent_context,
+            round_index=3,
+        )
+
+        result = context.to_dict()
+
+        assert isinstance(result, dict)
+        assert "message" in result
+        assert "sender" in result
+        assert "agent_context" in result
+        assert "round_index" in result
+        assert result["round_index"] == 3
+
+
+class TestAgentMessage:
+    """Test suite for AgentMessage class."""
+
+    def test_initialization_with_defaults(self):
+        """Test AgentMessage initialization with default values."""
+        message = AgentMessage()
+
+        assert message.content is None
+        assert message.name is None
+        assert message.rounds == 0
+        assert message.context is None
+        assert message.action_report is None
+        assert message.review_info is None
+        assert message.current_goal is None
+        assert message.model_name is None
+        assert message.role is None
+        assert message.success is True
+        assert message.resource_info is None
+
+    def test_initialization_with_custom_values(self):
+        """Test AgentMessage initialization with custom values."""
+        action_report = ActionOutput(content="Action result", is_exe_success=True)
+        review_info = AgentReviewInfo(approve=True, comments="Looks good")
+
+        message = AgentMessage(
+            content="Test content",
+            name="TestAgent",
+            rounds=3,
+            context={"key": "value"},
+            action_report=action_report,
+            review_info=review_info,
+            current_goal="Test goal",
+            model_name="gpt-4",
+            role=ModelMessageRoleType.AI,
+            success=False,
+            resource_info={"resource": "data"},
+        )
+
+        assert message.content == "Test content"
+        assert message.name == "TestAgent"
+        assert message.rounds == 3
+        assert message.context == {"key": "value"}
+        assert message.action_report == action_report
+        assert message.review_info == review_info
+        assert message.current_goal == "Test goal"
+        assert message.model_name == "gpt-4"
+        assert message.role == ModelMessageRoleType.AI
+        assert message.success is False
+        assert message.resource_info == {"resource": "data"}
+
+    def test_copy(self):
+        """Test AgentMessage.copy() method."""
+        action_report = ActionOutput(content="Action result")
+        review_info = AgentReviewInfo(approve=True, comments="Good")
+        original = AgentMessage(
+            content="Original content",
+            name="Agent1",
+            rounds=2,
+            context={"key": "value"},
+            action_report=action_report,
+            review_info=review_info,
+            current_goal="Goal",
+            model_name="gpt-3.5",
+            role=ModelMessageRoleType.HUMAN,
+            success=True,
+            resource_info={"res": "info"},
+        )
+
+        copied = original.copy()
+
+        # Check that values are copied
+        assert copied.content == original.content
+        assert copied.name == original.name
+        assert copied.rounds == original.rounds
+        assert copied.current_goal == original.current_goal
+        assert copied.model_name == original.model_name
+        assert copied.role == original.role
+        assert copied.success == original.success
+        assert copied.action_report == original.action_report
+        assert copied.resource_info == original.resource_info
+
+        # Check that context is deep copied (for dict)
+        assert copied.context == original.context
+        assert copied.context is not original.context  # Different objects
+
+        # Check review_info is copied
+        assert copied.review_info.approve == original.review_info.approve
+        assert copied.review_info.comments == original.review_info.comments
+
+    def test_copy_with_string_context(self):
+        """Test AgentMessage.copy() with string context."""
+        original = AgentMessage(content="Test", context="string context")
+        copied = original.copy()
+
+        assert copied.context == original.context
+        # String context is immutable, so it's the same object
+        assert copied.context is original.context
+
+    def test_to_dict(self):
+        """Test AgentMessage.to_dict() method."""
+        action_report = ActionOutput(
+            content="Action result", is_exe_success=True, view="chart"
+        )
+        message = AgentMessage(
+            content="Test content",
+            name="Agent1",
+            rounds=2,
+            action_report=action_report,
+            role=ModelMessageRoleType.AI,
+        )
+
+        result = message.to_dict()
+
+        assert isinstance(result, dict)
+        assert result["content"] == "Test content"
+        assert result["name"] == "Agent1"
+        assert result["rounds"] == 2
+        assert result["role"] == ModelMessageRoleType.AI
+        assert "action_report" in result
+        assert isinstance(result["action_report"], dict)
+        assert result["action_report"]["content"] == "Action result"
+        assert result["action_report"]["is_exe_success"] is True
+
+    def test_to_dict_without_action_report(self):
+        """Test AgentMessage.to_dict() without action_report."""
+        message = AgentMessage(content="Simple message", rounds=1)
+        result = message.to_dict()
+
+        assert result["content"] == "Simple message"
+        assert result["rounds"] == 1
+        assert result["action_report"] is None
+
+    def test_to_llm_message(self):
+        """Test AgentMessage.to_llm_message() method."""
+        message = AgentMessage(
+            content="Test content",
+            context={"key": "value"},
+            role=ModelMessageRoleType.HUMAN,
+            name="Agent1",
+            rounds=5,
+        )
+
+        result = message.to_llm_message()
+
+        assert isinstance(result, dict)
+        assert result["content"] == "Test content"
+        assert result["context"] == {"key": "value"}
+        assert result["role"] == ModelMessageRoleType.HUMAN
+        # to_llm_message only returns content, context, and role
+        assert "name" not in result
+        assert "rounds" not in result
+
+    def test_from_llm_message(self):
+        """Test AgentMessage.from_llm_message() factory method."""
+        llm_msg = {
+            "content": "Test content",
+            "context": {"key": "value"},
+            "role": ModelMessageRoleType.AI,
+            "rounds": 3,
+        }
+
+        message = AgentMessage.from_llm_message(llm_msg)
+
+        assert message.content == "Test content"
+        assert message.context == {"key": "value"}
+        assert message.role == ModelMessageRoleType.AI
+        assert message.rounds == 3
+
+    def test_from_llm_message_with_missing_fields(self):
+        """Test AgentMessage.from_llm_message() with missing fields."""
+        llm_msg = {"content": "Just content"}
+
+        message = AgentMessage.from_llm_message(llm_msg)
+
+        assert message.content == "Just content"
+        assert message.context is None
+        assert message.role is None
+        assert message.rounds == 0  # Default value
+
+    def test_from_messages(self):
+        """Test AgentMessage.from_messages() factory method."""
+        messages_data = [
+            {
+                "content": "Message 1",
+                "role": ModelMessageRoleType.HUMAN,
+                "rounds": 1,
+                "name": "User",
+            },
+            {
+                "content": "Message 2",
+                "role": ModelMessageRoleType.AI,
+                "rounds": 2,
+                "name": "Agent",
+                "success": False,
+            },
+            {"content": "Message 3"},
+        ]
+
+        messages = AgentMessage.from_messages(messages_data)
+
+        assert len(messages) == 3
+        assert messages[0].content == "Message 1"
+        assert messages[0].role == ModelMessageRoleType.HUMAN
+        assert messages[0].rounds == 1
+        assert messages[0].name == "User"
+        assert messages[1].content == "Message 2"
+        assert messages[1].role == ModelMessageRoleType.AI
+        assert messages[1].success is False
+        assert messages[2].content == "Message 3"
+
+    def test_from_messages_filters_unknown_fields(self):
+        """Test that from_messages filters out unknown fields."""
+        messages_data = [
+            {
+                "content": "Message 1",
+                "unknown_field": "should be ignored",
+                "another_unknown": 123,
+            }
+        ]
+
+        messages = AgentMessage.from_messages(messages_data)
+
+        assert len(messages) == 1
+        assert messages[0].content == "Message 1"
+        # Unknown fields should not cause errors
+
+    def test_action_report_serialization(self):
+        """Test action_report is properly serialized in to_dict()."""
+        action_report = ActionOutput(
+            content="Action completed",
+            is_exe_success=True,
+            thoughts="I thought about it",
+            observations="I observed this",
+            next_speakers=["Agent2", "Agent3"],
+        )
+        message = AgentMessage(content="Test", action_report=action_report)
+
+        result = message.to_dict()
+
+        assert isinstance(result["action_report"], dict)
+        assert result["action_report"]["content"] == "Action completed"
+        assert result["action_report"]["is_exe_success"] is True
+        assert result["action_report"]["thoughts"] == "I thought about it"
+        assert result["action_report"]["observations"] == "I observed this"
+        assert result["action_report"]["next_speakers"] == ["Agent2", "Agent3"]
+
+    def test_context_deep_copy_dict(self):
+        """Test that dict context is deep copied."""
+        original_context = {"key": "value", "nested": {"inner": "data"}}
+        original = AgentMessage(content="Test", context=original_context)
+
+        copied = original.copy()
+
+        # Modify the copied context
+        copied.context["key"] = "modified"
+        copied.context["nested"]["inner"] = "modified_data"
+
+        # Original should not be affected
+        assert original.context["key"] == "value"
+        assert original.context["nested"]["inner"] == "data"
+
+    def test_get_dict_context(self):
+        """Test AgentMessage.get_dict_context() method."""
+        # Test with dict context
+        message1 = AgentMessage(content="Test", context={"key": "value"})
+        result1 = message1.get_dict_context()
+        assert result1 == {"key": "value"}
+
+        # Test with string context
+        message2 = AgentMessage(content="Test", context="string context")
+        result2 = message2.get_dict_context()
+        assert result2 == {}
+
+        # Test with None context
+        message3 = AgentMessage(content="Test", context=None)
+        result3 = message3.get_dict_context()
+        assert result3 == {}
+
+
+class TestAgentReviewInfo:
+    """Test suite for AgentReviewInfo class."""
+
+    def test_initialization_with_defaults(self):
+        """Test AgentReviewInfo initialization with default values."""
+        review = AgentReviewInfo()
+
+        assert review.approve is False
+        assert review.comments is None
+
+    def test_initialization_with_custom_values(self):
+        """Test AgentReviewInfo initialization with custom values."""
+        review = AgentReviewInfo(approve=True, comments="Excellent work!")
+
+        assert review.approve is True
+        assert review.comments == "Excellent work!"
+
+    def test_copy(self):
+        """Test AgentReviewInfo.copy() method."""
+        original = AgentReviewInfo(approve=True, comments="Good job")
+        copied = original.copy()
+
+        assert copied.approve == original.approve
+        assert copied.comments == original.comments
+
+        # Modify copied
+        copied.approve = False
+        copied.comments = "Needs improvement"
+
+        # Original should not be affected
+        assert original.approve is True
+        assert original.comments == "Good job"
+
+    def test_to_dict(self):
+        """Test AgentReviewInfo.to_dict() method."""
+        review = AgentReviewInfo(approve=True, comments="Looks good")
+        result = review.to_dict()
+
+        assert isinstance(result, dict)
+        assert result["approve"] is True
+        assert result["comments"] == "Looks good"
+
+    def test_to_dict_with_defaults(self):
+        """Test AgentReviewInfo.to_dict() with default values."""
+        review = AgentReviewInfo()
+        result = review.to_dict()
+
+        assert isinstance(result, dict)
+        assert result["approve"] is False
+        assert result["comments"] is None
+
+
+@pytest.mark.parametrize(
+    "conv_id,max_chat_round,temperature",
+    [
+        ("conv_1", 50, 0.5),
+        ("conv_2", 100, 0.7),
+        ("conv_3", 200, 0.3),
+    ],
+)
+def test_agent_context_parametrized(conv_id, max_chat_round, temperature):
+    """Parametrized test for AgentContext with different values."""
+    context = AgentContext(
+        conv_id=conv_id, max_chat_round=max_chat_round, temperature=temperature
+    )
+
+    assert context.conv_id == conv_id
+    assert context.max_chat_round == max_chat_round
+    assert context.temperature == temperature
+
+
+@pytest.mark.parametrize(
+    "content,rounds,success",
+    [
+        ("Message 1", 1, True),
+        ("Message 2", 5, False),
+        ("Message 3", 10, True),
+    ],
+)
+def test_agent_message_parametrized(content, rounds, success):
+    """Parametrized test for AgentMessage with different values."""
+    message = AgentMessage(content=content, rounds=rounds, success=success)
+
+    assert message.content == content
+    assert message.rounds == rounds
+    assert message.success == success
--- a/tests/unit_tests/agent/test_awel_agent_operator.py
+++ b/tests/unit_tests/agent/test_awel_agent_operator.py
@@ -0,0 +1,746 @@
+"""Unit tests for AWEL agent operators.
+
+Tests for AWELAgentOperator, WrappedAgentOperator, and AgentBranchOperator
+from dbgpt.agent.core.plan.awel.agent_operator module.
+"""
+
+from unittest.mock import AsyncMock, MagicMock, patch
+
+import pytest
+
+from dbgpt.agent.core.action.base import ActionOutput
+from dbgpt.agent.core.agent import (
+    Agent,
+    AgentContext,
+    AgentGenerateContext,
+    AgentMessage,
+)
+from dbgpt.agent.core.base_agent import ConversableAgent
+from dbgpt.agent.core.memory.agent_memory import AgentMemory
+from dbgpt.agent.core.plan.awel.agent_operator import (
+    AgentBranchOperator,
+    AWELAgentOperator,
+    WrappedAgentOperator,
+)
+from dbgpt.agent.core.plan.awel.agent_operator_resource import AWELAgent
+from dbgpt.core.interface.message import ModelMessageRoleType
+
+
+@pytest.fixture
+def mock_agent():
+    """Create a mock Agent instance."""
+    agent = MagicMock(spec=Agent)
+    agent.name = "TestAgent"
+    agent.role = "test_role"
+    agent.send = AsyncMock()
+    agent.generate_reply = AsyncMock()
+    return agent
+
+
+@pytest.fixture
+def mock_conversable_agent():
+    """Create a mock ConversableAgent instance."""
+    agent = MagicMock(spec=ConversableAgent)
+    agent.name = "TestConversableAgent"
+    agent.role = "conversable_role"
+    agent.fixed_subgoal = None
+    agent.send = AsyncMock()
+    agent.generate_reply = AsyncMock()
+    return agent
+
+
+@pytest.fixture
+def mock_agent_memory():
+    """Create a mock AgentMemory instance."""
+    memory = MagicMock(spec=AgentMemory)
+    memory.structure_clone = MagicMock(return_value=memory)
+    return memory
+
+
+@pytest.fixture
+def sample_agent_context():
+    """Create a sample AgentContext."""
+    return AgentContext(conv_id="test_conv", max_chat_round=10, temperature=0.5)
+
+
+@pytest.fixture
+def sample_message():
+    """Create a sample AgentMessage."""
+    return AgentMessage(content="Test message content", rounds=0, success=True)
+
+
+@pytest.fixture
+def sample_generate_context(mock_agent, sample_message, sample_agent_context):
+    """Create a sample AgentGenerateContext."""
+    return AgentGenerateContext(
+        message=sample_message,
+        sender=mock_agent,
+        reviewer=None,
+        silent=False,
+        already_failed=False,
+        rely_messages=[],
+        agent_context=sample_agent_context,
+        llm_client=None,
+    )
+
+
+@pytest.fixture
+def mock_awel_agent():
+    """Create a mock AWELAgent instance."""
+    awel_agent = MagicMock(spec=AWELAgent)
+    awel_agent.agent_profile = "TestAgentProfile"
+    awel_agent.role_name = "TestRole"
+    awel_agent.llm_config = None
+    awel_agent.fixed_subgoal = None
+    awel_agent.agent_prompt = None
+    awel_agent.resources = []
+    return awel_agent
+
+
+class TestWrappedAgentOperator:
+    """Test suite for WrappedAgentOperator class."""
+
+    @pytest.mark.asyncio
+    async def test_map_basic_execution(self, mock_agent, sample_message):
+        """Test WrappedAgentOperator.map() with basic execution flow."""
+        # Setup
+        operator = WrappedAgentOperator(agent=mock_agent)
+        rely_msg = AgentMessage(
+            content="Previous message", role=ModelMessageRoleType.HUMAN
+        )
+        input_context = AgentGenerateContext(
+            message=sample_message,
+            sender=mock_agent,
+            rely_messages=[rely_msg],
+        )
+
+        # Mock the agent's generate_reply to return a successful message
+        reply_message = AgentMessage(
+            content="Agent reply", success=True, role=ModelMessageRoleType.AI
+        )
+        mock_agent.generate_reply.return_value = reply_message
+
+        # Execute
+        result = await operator.map(input_context)
+
+        # Verify
+        assert isinstance(result, AgentGenerateContext)
+        assert result.sender == mock_agent  # Sender is now the agent
+        assert len(result.rely_messages) == 2  # Human message + AI reply
+        assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN
+        assert result.rely_messages[1].role == ModelMessageRoleType.AI
+        assert result.already_started is True
+        assert result.last_speaker == mock_agent
+
+        # Verify agent.send was called
+        mock_agent.send.assert_called_once()
+
+        # Verify agent.generate_reply was called
+        mock_agent.generate_reply.assert_called_once()
+
+    @pytest.mark.asyncio
+    async def test_map_empty_message_raises_exception(self, mock_agent):
+        """Test WrappedAgentOperator.map() raises ValueError when message is empty."""
+        operator = WrappedAgentOperator(agent=mock_agent)
+        input_context = AgentGenerateContext(message=None, sender=mock_agent)
+
+        with pytest.raises(ValueError, match="The message is empty"):
+            await operator.map(input_context)
+
+    @pytest.mark.asyncio
+    async def test_map_empty_sender_raises_exception(self, mock_agent, sample_message):
+        """Test WrappedAgentOperator.map() raises ValueError when sender is empty."""
+        operator = WrappedAgentOperator(agent=mock_agent)
+        input_context = AgentGenerateContext(message=sample_message, sender=None)
+
+        with pytest.raises(ValueError, match="The sender is empty"):
+            await operator.map(input_context)
+
+    @pytest.mark.asyncio
+    async def test_map_unsuccessful_reply_raises_exception(
+        self, mock_agent, sample_message
+    ):
+        """Test WrappedAgentOperator.map() raises exception on unsuccessful reply."""
+        operator = WrappedAgentOperator(agent=mock_agent)
+        input_context = AgentGenerateContext(message=sample_message, sender=mock_agent)
+
+        # Mock the agent's generate_reply to return an unsuccessful message
+        reply_message = AgentMessage(
+            content="Failed to complete task",
+            success=False,
+            role=ModelMessageRoleType.AI,
+        )
+        mock_agent.generate_reply.return_value = reply_message
+
+        with pytest.raises(
+            ValueError,
+            match="The task failed at step test_role and the attempt to repair it",
+        ):
+            await operator.map(input_context)
+
+    @pytest.mark.asyncio
+    async def test_map_with_memory(self, mock_agent, sample_message, mock_agent_memory):
+        """Test WrappedAgentOperator.map() with memory handling."""
+        operator = WrappedAgentOperator(agent=mock_agent)
+        input_context = AgentGenerateContext(
+            message=sample_message, sender=mock_agent, memory=mock_agent_memory
+        )
+
+        reply_message = AgentMessage(content="Reply", success=True)
+        mock_agent.generate_reply.return_value = reply_message
+
+        result = await operator.map(input_context)
+
+        # Verify memory is cloned
+        mock_agent_memory.structure_clone.assert_called_once()
+        assert result.memory is not None
+
+    @pytest.mark.asyncio
+    async def test_map_current_goal_construction(self, mock_agent, sample_message):
+        """Test WrappedAgentOperator.map() constructs current_goal correctly."""
+        operator = WrappedAgentOperator(agent=mock_agent)
+        input_context = AgentGenerateContext(message=sample_message, sender=mock_agent)
+
+        reply_message = AgentMessage(content="Reply", success=True)
+        mock_agent.generate_reply.return_value = reply_message
+
+        await operator.map(input_context)
+
+        # Verify the message passed to generate_reply has current_goal set
+        call_args = mock_agent.generate_reply.call_args
+        received_message = call_args.kwargs["received_message"]
+        assert received_message.current_goal is not None
+        assert received_message.current_goal.startswith("[TestAgent]:")
+
+
+class TestAWELAgentOperator:
+    """Test suite for AWELAgentOperator class."""
+
+    @pytest.mark.asyncio
+    async def test_map_returns_early_if_already_failed(
+        self, mock_awel_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() returns early when already_failed is True."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+        sample_generate_context.already_failed = True
+
+        result = await operator.map(sample_generate_context)
+
+        assert result == sample_generate_context
+        assert result.already_failed is True
+
+    @pytest.mark.asyncio
+    async def test_map_empty_message_raises_exception(
+        self, mock_awel_agent, mock_agent, sample_agent_context
+    ):
+        """Test AWELAgentOperator.map() raises ValueError when message is empty."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+        input_context = AgentGenerateContext(
+            message=None, sender=mock_agent, agent_context=sample_agent_context
+        )
+
+        with pytest.raises(ValueError, match="The message is empty"):
+            await operator.map(input_context)
+
+    @pytest.mark.asyncio
+    async def test_map_empty_sender_raises_exception(
+        self, mock_awel_agent, sample_message, sample_agent_context
+    ):
+        """Test AWELAgentOperator.map() raises ValueError when sender is empty."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+        input_context = AgentGenerateContext(
+            message=sample_message, sender=None, agent_context=sample_agent_context
+        )
+
+        # Need to mock get_agent to avoid issues
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = MagicMock(spec=ConversableAgent)
+
+            with pytest.raises(ValueError, match="The sender is empty"):
+                await operator.map(input_context)
+
+    @pytest.mark.asyncio
+    async def test_map_normal_execution(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() normal execution flow."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        # Mock get_agent to return our conversable agent
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            # Mock agent's generate_reply
+            reply_message = AgentMessage(
+                content="Agent response", success=True, role=ModelMessageRoleType.AI
+            )
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(sample_generate_context)
+
+            # Verify
+            assert isinstance(result, AgentGenerateContext)
+            assert result.sender == mock_conversable_agent
+            assert len(result.rely_messages) == 2  # HUMAN + AI messages
+            assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN
+            assert result.rely_messages[1].role == ModelMessageRoleType.AI
+            assert result.already_started is True
+            assert result.already_failed is False
+            assert result.last_speaker == mock_conversable_agent
+
+    @pytest.mark.asyncio
+    async def test_map_with_fixed_subgoal(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() with fixed_subgoal."""
+        mock_conversable_agent.fixed_subgoal = "Fixed goal for testing"
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            await operator.map(sample_generate_context)
+
+            # Verify that fixed_subgoal overrides message content
+            call_args = mock_conversable_agent.generate_reply.call_args
+            received_message = call_args.kwargs["received_message"]
+            assert received_message.content == "Fixed goal for testing"
+            assert "Fixed goal for testing" in received_message.current_goal
+
+    @pytest.mark.asyncio
+    async def test_map_without_fixed_subgoal(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() without fixed_subgoal."""
+        mock_conversable_agent.fixed_subgoal = None
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            await operator.map(sample_generate_context)
+
+            # Verify that message content is used
+            call_args = mock_conversable_agent.generate_reply.call_args
+            received_message = call_args.kwargs["received_message"]
+            assert received_message.content == "Test message content"
+
+    @pytest.mark.asyncio
+    async def test_map_failure_sets_already_failed(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() sets already_failed on unsuccessful reply."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            # Mock agent's generate_reply to return unsuccessful message
+            reply_message = AgentMessage(
+                content="Failed", success=False, role=ModelMessageRoleType.AI
+            )
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(sample_generate_context)
+
+            # Verify already_failed is set to True
+            assert result.already_failed is True
+
+    @pytest.mark.asyncio
+    async def test_map_increments_rounds(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() increments message rounds."""
+        sample_generate_context.message.rounds = 5
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            await operator.map(sample_generate_context)
+
+            # Verify rounds are incremented
+            call_args = mock_conversable_agent.generate_reply.call_args
+            received_message = call_args.kwargs["received_message"]
+            assert received_message.rounds == 6
+
+    @pytest.mark.asyncio
+    async def test_map_with_begin_agent_not_started(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() when begin_agent is set but not started."""
+        sample_generate_context.begin_agent = "other_agent"
+        sample_generate_context.already_started = False
+        mock_conversable_agent.role = "test_role"
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            result = await operator.map(sample_generate_context)
+
+            # Should return early without executing
+            assert result == sample_generate_context
+            mock_conversable_agent.generate_reply.assert_not_called()
+
+    @pytest.mark.asyncio
+    async def test_map_with_begin_agent_matching(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() when begin_agent matches current agent."""
+        sample_generate_context.begin_agent = "conversable_role"
+        sample_generate_context.already_started = False
+        mock_conversable_agent.role = "conversable_role"
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(sample_generate_context)
+
+            # Should execute and set is_retry_chat=True
+            assert result.sender == mock_conversable_agent
+            mock_conversable_agent.generate_reply.assert_called_once()
+            call_args = mock_conversable_agent.generate_reply.call_args
+            assert call_args.kwargs.get("is_retry_chat") is True
+
+    @pytest.mark.asyncio
+    async def test_map_with_memory_clone(
+        self,
+        mock_awel_agent,
+        mock_conversable_agent,
+        sample_generate_context,
+        mock_agent_memory,
+    ):
+        """Test AWELAgentOperator.map() clones memory."""
+        sample_generate_context.memory = mock_agent_memory
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(sample_generate_context)
+
+            # Verify memory is cloned
+            mock_agent_memory.structure_clone.assert_called_once()
+            assert result.memory is not None
+
+    @pytest.mark.asyncio
+    async def test_map_message_role_conversion(
+        self, mock_awel_agent, mock_conversable_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.map() converts message roles correctly."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(sample_generate_context)
+
+            # Check that roles are set correctly
+            assert result.rely_messages[0].role == ModelMessageRoleType.HUMAN
+            assert result.rely_messages[1].role == ModelMessageRoleType.AI
+
+    @pytest.mark.asyncio
+    async def test_get_agent_builds_agent(
+        self, mock_awel_agent, sample_generate_context
+    ):
+        """Test AWELAgentOperator.get_agent() builds agent correctly."""
+        operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+        operator.llm_client = MagicMock()
+
+        # Mock the agent class and its builder pattern
+        mock_agent_class = MagicMock()
+        mock_agent_instance = MagicMock(spec=ConversableAgent)
+        mock_agent_instance.role = "test_role"
+
+        # Setup the builder chain
+        mock_builder = MagicMock()
+        mock_builder.bind.return_value = mock_builder
+        mock_builder.build = AsyncMock(return_value=mock_agent_instance)
+        mock_agent_class.return_value = mock_builder
+
+        with patch(
+            "dbgpt.agent.core.plan.awel.agent_operator.get_agent_manager"
+        ) as mock_get_manager:
+            mock_manager = MagicMock()
+            mock_manager.get_by_name.return_value = mock_agent_class
+            mock_get_manager.return_value = mock_manager
+
+            with patch(
+                "dbgpt.agent.core.plan.awel.agent_operator.get_resource_manager"
+            ) as mock_get_resource_manager:
+                mock_resource_manager = MagicMock()
+                mock_resource_manager.build_resource.return_value = None
+                mock_get_resource_manager.return_value = mock_resource_manager
+
+                result = await operator.get_agent(sample_generate_context)
+
+                assert result == mock_agent_instance
+                mock_manager.get_by_name.assert_called_once_with("TestAgentProfile")
+
+
+class TestAgentBranchOperator:
+    """Test suite for AgentBranchOperator class."""
+
+    @pytest.mark.asyncio
+    async def test_branches_returns_branch_map(self):
+        """Test AgentBranchOperator.branches() returns correct branch map."""
+        operator = AgentBranchOperator()
+
+        # Create mock downstream AWELAgentOperator nodes
+        mock_awel_agent1 = MagicMock(spec=AWELAgent)
+        mock_awel_agent1.agent_profile = "Agent1"
+
+        mock_awel_agent2 = MagicMock(spec=AWELAgent)
+        mock_awel_agent2.agent_profile = "Agent2"
+
+        mock_operator1 = MagicMock(spec=AWELAgentOperator)
+        mock_operator1.awel_agent = mock_awel_agent1
+        mock_operator1.node_name = "node1"
+
+        mock_operator2 = MagicMock(spec=AWELAgentOperator)
+        mock_operator2.awel_agent = mock_awel_agent2
+        mock_operator2.node_name = "node2"
+
+        # Mock downstream attribute
+        operator.downstream = [mock_operator1, mock_operator2]
+
+        result = await operator.branches()
+
+        assert isinstance(result, dict)
+        assert len(result) == 2
+
+        # The keys should be functions and values should be node names
+        list(result.keys())
+        branch_targets = list(result.values())
+        assert branch_targets == ["node1", "node2"]
+
+    @pytest.mark.asyncio
+    async def test_branch_function_checks_next_speakers(self):
+        """Test that branch functions check action_report.next_speakers."""
+        operator = AgentBranchOperator()
+
+        mock_awel_agent = MagicMock(spec=AWELAgent)
+        mock_awel_agent.agent_profile = "TargetAgent"
+
+        mock_operator = MagicMock(spec=AWELAgentOperator)
+        mock_operator.awel_agent = mock_awel_agent
+        mock_operator.node_name = "target_node"
+
+        operator.downstream = [mock_operator]
+
+        branch_map = await operator.branches()
+        branch_func = list(branch_map.keys())[0]
+
+        # Test with matching next_speakers
+        action_output = ActionOutput(
+            content="Test", next_speakers=["TargetAgent", "OtherAgent"]
+        )
+        last_message = AgentMessage(content="Test", action_report=action_output)
+        context = AgentGenerateContext(
+            message=AgentMessage(content="Main"),
+            sender=MagicMock(spec=Agent),
+            rely_messages=[last_message],
+        )
+
+        result = branch_func(context)
+        assert result is True
+
+        # Test with non-matching next_speakers
+        action_output2 = ActionOutput(content="Test", next_speakers=["OtherAgent"])
+        last_message2 = AgentMessage(content="Test", action_report=action_output2)
+        context2 = AgentGenerateContext(
+            message=AgentMessage(content="Main"),
+            sender=MagicMock(spec=Agent),
+            rely_messages=[last_message2],
+        )
+
+        result2 = branch_func(context2)
+        assert result2 is False
+
+    @pytest.mark.asyncio
+    async def test_branch_function_returns_false_without_action_report(self):
+        """Test branch function returns False when action_report is None."""
+        operator = AgentBranchOperator()
+
+        mock_awel_agent = MagicMock(spec=AWELAgent)
+        mock_awel_agent.agent_profile = "TargetAgent"
+
+        mock_operator = MagicMock(spec=AWELAgentOperator)
+        mock_operator.awel_agent = mock_awel_agent
+        mock_operator.node_name = "target_node"
+
+        operator.downstream = [mock_operator]
+
+        branch_map = await operator.branches()
+        branch_func = list(branch_map.keys())[0]
+
+        # Test without action_report
+        last_message = AgentMessage(content="Test", action_report=None)
+        context = AgentGenerateContext(
+            message=AgentMessage(content="Main"),
+            sender=MagicMock(spec=Agent),
+            rely_messages=[last_message],
+        )
+
+        result = branch_func(context)
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_branch_function_returns_false_without_next_speakers(self):
+        """Test branch function returns False when next_speakers is None."""
+        operator = AgentBranchOperator()
+
+        mock_awel_agent = MagicMock(spec=AWELAgent)
+        mock_awel_agent.agent_profile = "TargetAgent"
+
+        mock_operator = MagicMock(spec=AWELAgentOperator)
+        mock_operator.awel_agent = mock_awel_agent
+        mock_operator.node_name = "target_node"
+
+        operator.downstream = [mock_operator]
+
+        branch_map = await operator.branches()
+        branch_func = list(branch_map.keys())[0]
+
+        # Test with action_report but no next_speakers
+        action_output = ActionOutput(content="Test", next_speakers=None)
+        last_message = AgentMessage(content="Test", action_report=action_output)
+        context = AgentGenerateContext(
+            message=AgentMessage(content="Main"),
+            sender=MagicMock(spec=Agent),
+            rely_messages=[last_message],
+        )
+
+        result = branch_func(context)
+        assert result is False
+
+    @pytest.mark.asyncio
+    async def test_branches_filters_non_awel_operators(self):
+        """Test branches() filters out non-AWELAgentOperator nodes."""
+        operator = AgentBranchOperator()
+
+        mock_awel_agent = MagicMock(spec=AWELAgent)
+        mock_awel_agent.agent_profile = "Agent1"
+
+        mock_operator1 = MagicMock(spec=AWELAgentOperator)
+        mock_operator1.awel_agent = mock_awel_agent
+        mock_operator1.node_name = "node1"
+
+        # Add a non-AWELAgentOperator node
+        mock_other_operator = MagicMock()
+
+        operator.downstream = [mock_operator1, mock_other_operator]
+
+        result = await operator.branches()
+
+        # Should only include the AWELAgentOperator
+        assert len(result) == 1
+
+
+@pytest.mark.parametrize(
+    "already_failed,should_execute",
+    [
+        (True, False),  # Should return early
+        (False, True),  # Should execute
+    ],
+)
+@pytest.mark.asyncio
+async def test_awel_operator_already_failed_parametrized(
+    already_failed, should_execute, mock_awel_agent, mock_conversable_agent
+):
+    """Parametrized test for AWELAgentOperator with already_failed flag."""
+    operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+    message = AgentMessage(content="Test", rounds=0)
+    mock_sender = MagicMock(spec=Agent)
+
+    input_context = AgentGenerateContext(
+        message=message,
+        sender=mock_sender,
+        already_failed=already_failed,
+        agent_context=AgentContext(conv_id="test"),
+    )
+
+    if should_execute:
+        with patch.object(
+            operator, "get_agent", new_callable=AsyncMock
+        ) as mock_get_agent:
+            mock_get_agent.return_value = mock_conversable_agent
+            reply_message = AgentMessage(content="Response", success=True)
+            mock_conversable_agent.generate_reply.return_value = reply_message
+
+            result = await operator.map(input_context)
+            assert result.sender == mock_conversable_agent
+    else:
+        result = await operator.map(input_context)
+        assert result == input_context
+
+
+@pytest.mark.parametrize(
+    "success,expected_already_failed",
+    [
+        (True, False),
+        (False, True),
+    ],
+)
+@pytest.mark.asyncio
+async def test_awel_operator_success_flag_parametrized(
+    success, expected_already_failed, mock_awel_agent, mock_conversable_agent
+):
+    """Parametrized test for AWELAgentOperator success handling."""
+    operator = AWELAgentOperator(awel_agent=mock_awel_agent)
+    message = AgentMessage(content="Test", rounds=0)
+    mock_sender = MagicMock(spec=Agent)
+
+    input_context = AgentGenerateContext(
+        message=message,
+        sender=mock_sender,
+        agent_context=AgentContext(conv_id="test"),
+    )
+
+    with patch.object(operator, "get_agent", new_callable=AsyncMock) as mock_get_agent:
+        mock_get_agent.return_value = mock_conversable_agent
+        reply_message = AgentMessage(content="Response", success=success)
+        mock_conversable_agent.generate_reply.return_value = reply_message
+
+        result = await operator.map(input_context)
+        assert result.already_failed == expected_already_failed