update tests

python/packages/autogen-agentchat/tests/test_declarative_components.py

@@ -1,7 +1,15 @@
 import pytest
-from autogen_agentchat.base import OrTerminationCondition
-from autogen_agentchat.conditions import MaxMessageTermination, StopMessageTermination
-from autogen_agentchat.messages import StopMessage
+from autogen_agentchat.base import AndTerminationCondition
+from autogen_agentchat.conditions import (
+    ExternalTermination,
+    HandoffTermination,
+    MaxMessageTermination,
+    SourceMatchTermination,
+    StopMessageTermination,
+    TextMentionTermination,
+    TimeoutTermination,
+    TokenUsageTermination,
+)
 from autogen_core import ComponentLoader, ComponentModel
 
 
@@ -11,41 +19,76 @@ async def test_termination_declarative() -> None:
     # Create basic termination conditions
     max_term = MaxMessageTermination(5)
     stop_term = StopMessageTermination()
+    text_term = TextMentionTermination("stop")
+    token_term = TokenUsageTermination(max_total_token=100, max_prompt_token=50, max_completion_token=100)
+    handoff_term = HandoffTermination(target="human")
+    timeout_term = TimeoutTermination(timeout_seconds=30)
+    external_term = ExternalTermination()
+    source_term = SourceMatchTermination(sources=["human"])
 
     # Test basic serialization
     max_config = max_term.dump_component()
     assert isinstance(max_config, ComponentModel)
     assert max_config.provider == "autogen_agentchat.conditions.MaxMessageTermination"
     assert max_config.config.get("max_messages") == 5
 
+    # Test serialization of new conditions
+    text_config = text_term.dump_component()
+    assert text_config.provider == "autogen_agentchat.conditions.TextMentionTermination"
+    assert text_config.config.get("text") == "stop"
+
+    token_config = token_term.dump_component()
+    assert token_config.provider == "autogen_agentchat.conditions.TokenUsageTermination"
+    assert token_config.config.get("max_total_token") == 100
+
+    handoff_config = handoff_term.dump_component()
+    assert handoff_config.provider == "autogen_agentchat.conditions.HandoffTermination"
+    assert handoff_config.config.get("target") == "human"
+
+    timeout_config = timeout_term.dump_component()
+    assert timeout_config.provider == "autogen_agentchat.conditions.TimeoutTermination"
+    assert timeout_config.config.get("timeout_seconds") == 30
+
+    external_config = external_term.dump_component()
+    assert external_config.provider == "autogen_agentchat.conditions.ExternalTermination"
+
+    source_config = source_term.dump_component()
+    assert source_config.provider == "autogen_agentchat.conditions.SourceMatchTermination"
+    assert source_config.config.get("sources") == ["human"]
+
     # Test basic deserialization
     loaded_max = ComponentLoader.load_component(max_config, MaxMessageTermination)
     assert isinstance(loaded_max, MaxMessageTermination)
-    # Use public interface to verify state
-    messages = [StopMessage(content="msg", source="test") for _ in range(5)]
-    result = await loaded_max(messages)
-    assert isinstance(result, StopMessage)
-
-    # Test composition and complex serialization
-    or_term = max_term | stop_term
-    or_config = or_term.dump_component()
-    assert or_config.provider == "autogen_agentchat.base.OrTerminationCondition"
-    assert len(or_config.config["conditions"]) == 2
-
-    # Verify nested conditions are correctly serialized
-    conditions = or_config.config["conditions"]
-    assert conditions[0]["provider"] == "autogen_agentchat.conditions.MaxMessageTermination"
-    assert conditions[1]["provider"] == "autogen_agentchat.conditions.StopMessageTermination"
-
-    # Test behavior of loaded composite condition
-    loaded_or = OrTerminationCondition.load_component(or_config)
-
-    # Test with stop message
-    stop_messages = [StopMessage(content="stop", source="test")]
-    result = await loaded_or(stop_messages)
-    assert isinstance(result, StopMessage)
-    assert loaded_or.terminated
-
-    # Test reset functionality
-    await loaded_or.reset()
-    assert not loaded_or.terminated
+
+    # Test deserialization of new conditions
+    loaded_text = ComponentLoader.load_component(text_config, TextMentionTermination)
+    assert isinstance(loaded_text, TextMentionTermination)
+
+    loaded_token = ComponentLoader.load_component(token_config, TokenUsageTermination)
+    assert isinstance(loaded_token, TokenUsageTermination)
+
+    loaded_handoff = ComponentLoader.load_component(handoff_config, HandoffTermination)
+    assert isinstance(loaded_handoff, HandoffTermination)
+
+    loaded_timeout = ComponentLoader.load_component(timeout_config, TimeoutTermination)
+    assert isinstance(loaded_timeout, TimeoutTermination)
+
+    loaded_external = ComponentLoader.load_component(external_config, ExternalTermination)
+    assert isinstance(loaded_external, ExternalTermination)
+
+    loaded_source = ComponentLoader.load_component(source_config, SourceMatchTermination)
+    assert isinstance(loaded_source, SourceMatchTermination)
+
+    # Test composition with new conditions
+    composite_term = (max_term | stop_term) & (token_term | handoff_term)
+    composite_config = composite_term.dump_component()
+
+    assert composite_config.provider == "autogen_agentchat.base.AndTerminationCondition"
+    conditions = composite_config.config["conditions"]
+    assert len(conditions) == 2
+    assert conditions[0]["provider"] == "autogen_agentchat.base.OrTerminationCondition"
+    assert conditions[1]["provider"] == "autogen_agentchat.base.OrTerminationCondition"
+
+    # Test loading complex composition
+    loaded_composite = ComponentLoader.load_component(composite_config)
+    assert isinstance(loaded_composite, AndTerminationCondition)