test/agent/test_context_variables.py
Rama main · réplica de lectura del repositorio original.
# Copyright (c) 2026, AG2ai, Inc., AG2ai open-source projects maintainers and core contributors
#
# SPDX-License-Identifier: Apache-2.0
import logging
from typing import Annotated, Any
from unittest.mock import MagicMock
import pytest
from ag2 import Agent, Context, Inject, Variable
from ag2.context import strip_reserved_variables
from ag2.events import ToolCallEvent
from ag2.middleware.builtin.tools.approval import BYPASS_KEY
from ag2.testing import TestConfig
@pytest.fixture()
def test_config() -> TestConfig:
return TestConfig(
ToolCallEvent(name="my_tool"),
"result",
)
@pytest.mark.asyncio()
async def test_ask_variables(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(ctx: Context) -> str:
mock(ctx.variables["dep"])
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
)
await agent.ask("Hi!", variables={"dep": "1"})
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_agent_variables(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(ctx: Context) -> str:
mock(ctx.variables["dep"])
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
variables={"dep": "1"},
)
await agent.ask("Hi!")
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_mixed_variables(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(ctx: Context) -> str:
mock(ctx.variables)
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
variables={"dep": "1"},
)
await agent.ask("Hi!", variables={"dep2": "2"})
mock.assert_called_once_with({"dep": "1", "dep2": "2"})
@pytest.mark.asyncio()
async def test_variable_alias(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(dep: Annotated[str, Variable()]) -> str:
mock(dep)
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
variables={"dep": "1"},
)
await agent.ask("Hi!")
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_variable_by_name(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(d: Annotated[str, Variable("dep")]) -> str:
mock(d)
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
variables={"dep": "1"},
)
await agent.ask("Hi!")
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_variable_with_default(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(dep: Annotated[str, Variable(default="1")]) -> str:
mock(dep)
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
)
await agent.ask("Hi!")
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_variable_with_default_factory(
mock: MagicMock,
test_config: TestConfig,
) -> None:
def my_tool(dep: Annotated[str, Variable(default_factory=dict)]) -> str:
mock(dep)
return ""
agent = Agent(
"",
config=test_config,
tools=[my_tool],
)
await agent.ask("Hi!")
mock.assert_called_once_with({})
@pytest.mark.asyncio()
async def test_set_variable_by_tool(mock: MagicMock) -> None:
def my_tool(ctx: Context) -> str:
assert not ctx.variables
ctx.variables["dep"] = "1"
return ""
def another_tool(ctx: Context) -> str:
mock(ctx.variables["dep"])
return ""
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
ToolCallEvent(name="another_tool"),
"result",
),
tools=[my_tool, another_tool],
)
await agent.ask("Hi!")
mock.assert_called_once_with("1")
@pytest.mark.asyncio()
async def test_variable_with_default_factory_called_once(mock: MagicMock) -> None:
def factory() -> list[int]:
mock.factory()
return [1]
def my_tool(
dep: Annotated[list[int], Variable(default_factory=factory)],
) -> str:
mock.first(dep.copy())
dep.append(2)
return ""
def another_tool(
dep: Annotated[list[int], Variable(default_factory=factory)],
) -> str:
mock.second(dep.copy())
return ""
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
ToolCallEvent(name="another_tool"),
"result",
),
tools=[my_tool, another_tool],
)
await agent.ask("Hi!")
mock.factory.assert_called_once()
mock.first.assert_called_once_with([1])
mock.second.assert_called_once_with([1, 2])
@pytest.mark.asyncio()
@pytest.mark.parametrize("value", ["", 0, False, [], {}], ids=repr)
async def test_falsy_values_are_resolved(mock: MagicMock, value: Any) -> None:
def my_tool(
var: Annotated[Any, Variable("myvar")],
dep: Annotated[Any, Inject("mydep")],
) -> str:
mock(var, dep)
return ""
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
variables={"myvar": value},
dependencies={"mydep": value},
)
await agent.ask("Hi!")
mock.assert_called_once_with(value, value)
@pytest.mark.asyncio()
@pytest.mark.parametrize("value", ["", 0, False, [], {}], ids=repr)
async def test_falsy_values_win_over_defaults(mock: MagicMock, value: Any) -> None:
def my_tool(
ctx: Context,
var: Annotated[Any, Variable("myvar", default="default")],
dep: Annotated[Any, Inject("mydep", default="default")],
) -> str:
mock(var, dep, ctx.variables["myvar"], ctx.dependencies["mydep"])
return ""
agent = Agent(
"",
config=TestConfig(
ToolCallEvent(name="my_tool"),
"result",
),
tools=[my_tool],
variables={"myvar": value},
dependencies={"mydep": value},
)
await agent.ask("Hi!")
mock.assert_called_once_with(value, value, value, value)
class TestReservedVariables:
"""Keys under ``ag:``/``a2a:`` are the framework's own; a peer never authors them."""
def test_ordinary_keys_survive(self) -> None:
payload = {"city": "Tokyo", "count": 3}
assert strip_reserved_variables(payload, source="a test") == payload
def test_reserved_keys_are_dropped(self) -> None:
payload = {"city": "Tokyo", BYPASS_KEY: {"pay": True}, "a2a:tenant": "acme"}
assert strip_reserved_variables(payload, source="a test") == {"city": "Tokyo"}
def test_a_drop_is_logged(self, caplog: pytest.LogCaptureFixture) -> None:
with caplog.at_level(logging.WARNING, logger="ag2.context"):
strip_reserved_variables({BYPASS_KEY: {"pay": True}}, source="a hostile peer")
assert BYPASS_KEY in caplog.text
assert "a hostile peer" in caplog.text
def test_the_outbound_side_drops_quietly(self, caplog: pytest.LogCaptureFixture) -> None:
with caplog.at_level(logging.WARNING, logger="ag2.context"):
kept = strip_reserved_variables({BYPASS_KEY: {"pay": True}}, source="a response", warn=False)
assert kept == {}
assert caplog.text == ""
def test_a_non_string_key_is_kept(self) -> None:
payload: dict[Any, Any] = {7: "seven"}
assert strip_reserved_variables(payload, source="a test") == payload