mirror of
https://github.com/Significant-Gravitas/AutoGPT.git
synced 2026-04-08 03:00:28 -04:00
a8fe383481
- Revert get_credits → get_credit_balance rename in DatabaseManager and DatabaseManagerClient to preserve existing RPC route name - Keep DatabaseManagerAsyncClient credits additions (new, no route change) - Remove text_utils.py (belongs to PR #12400) - Revert sidebar.tsx change (unrelated) - Restore removed comments in coercion tests
513 lines
16 KiB
Python
513 lines
16 KiB
Python
"""Tests for execute_block — credit charging and type coercion."""
|
|
|
|
from collections.abc import AsyncIterator
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
from backend.blocks._base import BlockType
|
|
from backend.copilot.tools.helpers import execute_block
|
|
from backend.copilot.tools.models import BlockOutputResponse, ErrorResponse
|
|
|
|
_USER = "test-user-helpers"
|
|
_SESSION = "test-session-helpers"
|
|
|
|
|
|
def _make_block(block_id: str = "block-1", name: str = "TestBlock"):
|
|
"""Create a minimal mock block for execute_block()."""
|
|
mock = MagicMock()
|
|
mock.id = block_id
|
|
mock.name = name
|
|
mock.block_type = BlockType.STANDARD
|
|
|
|
mock.input_schema = MagicMock()
|
|
mock.input_schema.get_credentials_fields_info.return_value = {}
|
|
|
|
async def _execute(
|
|
input_data: dict, **kwargs: Any
|
|
) -> AsyncIterator[tuple[str, Any]]:
|
|
yield "result", "ok"
|
|
|
|
mock.execute = _execute
|
|
return mock
|
|
|
|
|
|
def _patch_workspace():
|
|
"""Patch workspace_db to return a mock workspace."""
|
|
mock_workspace = MagicMock()
|
|
mock_workspace.id = "ws-1"
|
|
mock_ws_db = MagicMock()
|
|
mock_ws_db.get_or_create_workspace = AsyncMock(return_value=mock_workspace)
|
|
return patch("backend.copilot.tools.helpers.workspace_db", return_value=mock_ws_db)
|
|
|
|
|
|
def _patch_credit_db(
|
|
get_credits_return: int = 100,
|
|
spend_credits_side_effect: Any = None,
|
|
):
|
|
"""Patch credit_db accessor to return a mock credit adapter."""
|
|
mock_credit = MagicMock()
|
|
mock_credit.get_credits = AsyncMock(return_value=get_credits_return)
|
|
if spend_credits_side_effect is not None:
|
|
mock_credit.spend_credits = AsyncMock(side_effect=spend_credits_side_effect)
|
|
else:
|
|
mock_credit.spend_credits = AsyncMock()
|
|
return (
|
|
patch(
|
|
"backend.copilot.tools.helpers.credit_db",
|
|
return_value=mock_credit,
|
|
),
|
|
mock_credit,
|
|
)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Credit charging tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
class TestExecuteBlockCreditCharging:
|
|
async def test_charges_credits_when_cost_is_positive(self):
|
|
"""Block with cost > 0 should call spend_credits after execution."""
|
|
block = _make_block()
|
|
credit_patch, mock_credit = _patch_credit_db(get_credits_return=100)
|
|
|
|
with (
|
|
_patch_workspace(),
|
|
patch(
|
|
"backend.copilot.tools.helpers.block_usage_cost",
|
|
return_value=(10, {"key": "val"}),
|
|
),
|
|
credit_patch,
|
|
):
|
|
result = await execute_block(
|
|
block=block,
|
|
block_id="block-1",
|
|
input_data={"text": "hello"},
|
|
user_id=_USER,
|
|
session_id=_SESSION,
|
|
node_exec_id="exec-1",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(result, BlockOutputResponse)
|
|
assert result.success is True
|
|
mock_credit.spend_credits.assert_awaited_once()
|
|
call_kwargs = mock_credit.spend_credits.call_args.kwargs
|
|
assert call_kwargs["cost"] == 10
|
|
assert call_kwargs["metadata"].reason == "copilot_block_execution"
|
|
|
|
async def test_returns_error_when_insufficient_credits_before_exec(self):
|
|
"""Pre-execution check should return ErrorResponse when balance < cost."""
|
|
block = _make_block()
|
|
credit_patch, mock_credit = _patch_credit_db(get_credits_return=5)
|
|
|
|
with (
|
|
_patch_workspace(),
|
|
patch(
|
|
"backend.copilot.tools.helpers.block_usage_cost",
|
|
return_value=(10, {}),
|
|
),
|
|
credit_patch,
|
|
):
|
|
result = await execute_block(
|
|
block=block,
|
|
block_id="block-1",
|
|
input_data={},
|
|
user_id=_USER,
|
|
session_id=_SESSION,
|
|
node_exec_id="exec-1",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(result, ErrorResponse)
|
|
assert "Insufficient credits" in result.message
|
|
|
|
async def test_no_charge_when_cost_is_zero(self):
|
|
"""Block with cost 0 should not call spend_credits."""
|
|
block = _make_block()
|
|
credit_patch, mock_credit = _patch_credit_db()
|
|
|
|
with (
|
|
_patch_workspace(),
|
|
patch(
|
|
"backend.copilot.tools.helpers.block_usage_cost",
|
|
return_value=(0, {}),
|
|
),
|
|
credit_patch,
|
|
):
|
|
result = await execute_block(
|
|
block=block,
|
|
block_id="block-1",
|
|
input_data={},
|
|
user_id=_USER,
|
|
session_id=_SESSION,
|
|
node_exec_id="exec-1",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(result, BlockOutputResponse)
|
|
assert result.success is True
|
|
# Credit functions should not be called at all for zero-cost blocks
|
|
mock_credit.get_credits.assert_not_awaited()
|
|
mock_credit.spend_credits.assert_not_awaited()
|
|
|
|
async def test_returns_output_on_post_exec_insufficient_balance(self):
|
|
"""If charging fails after execution, output is still returned (block already ran)."""
|
|
from backend.util.exceptions import InsufficientBalanceError
|
|
|
|
block = _make_block()
|
|
credit_patch, mock_credit = _patch_credit_db(
|
|
get_credits_return=15,
|
|
spend_credits_side_effect=InsufficientBalanceError(
|
|
"Low balance", _USER, 5, 10
|
|
),
|
|
)
|
|
|
|
with (
|
|
_patch_workspace(),
|
|
patch(
|
|
"backend.copilot.tools.helpers.block_usage_cost",
|
|
return_value=(10, {}),
|
|
),
|
|
credit_patch,
|
|
):
|
|
result = await execute_block(
|
|
block=block,
|
|
block_id="block-1",
|
|
input_data={},
|
|
user_id=_USER,
|
|
session_id=_SESSION,
|
|
node_exec_id="exec-1",
|
|
matched_credentials={},
|
|
)
|
|
|
|
# Block already executed (with side effects), so output is returned
|
|
assert isinstance(result, BlockOutputResponse)
|
|
assert result.success is True
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Type coercion tests
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _make_block_schema(annotations: dict[str, Any]) -> MagicMock:
|
|
"""Create a mock input_schema with model_fields matching the given annotations."""
|
|
schema = MagicMock()
|
|
# coerce_inputs_to_schema uses model_fields (Pydantic v2 API)
|
|
model_fields = {}
|
|
for name, ann in annotations.items():
|
|
field = MagicMock()
|
|
field.annotation = ann
|
|
model_fields[name] = field
|
|
schema.model_fields = model_fields
|
|
return schema
|
|
|
|
|
|
def _make_coerce_block(
|
|
block_id: str,
|
|
name: str,
|
|
annotations: dict[str, Any],
|
|
outputs: dict[str, list[Any]] | None = None,
|
|
) -> MagicMock:
|
|
"""Create a mock block with typed annotations and a simple execute method."""
|
|
block = MagicMock()
|
|
block.id = block_id
|
|
block.name = name
|
|
block.input_schema = _make_block_schema(annotations)
|
|
|
|
captured_inputs: dict[str, Any] = {}
|
|
|
|
async def mock_execute(input_data: dict, **_kwargs: Any):
|
|
captured_inputs.update(input_data)
|
|
for output_name, values in (outputs or {"result": ["ok"]}).items():
|
|
for v in values:
|
|
yield output_name, v
|
|
|
|
block.execute = mock_execute
|
|
block._captured_inputs = captured_inputs
|
|
return block
|
|
|
|
|
|
_TEST_SESSION_ID = "test-session-coerce"
|
|
_TEST_USER_ID = "test-user-coerce"
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_json_string_to_nested_list():
|
|
"""JSON string → list[list[str]] (Google Sheets CSV import case)."""
|
|
block = _make_coerce_block(
|
|
"sheets-write",
|
|
"Google Sheets Write",
|
|
{"values": list[list[str]], "spreadsheet_id": str},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="sheets-write",
|
|
input_data={
|
|
"values": '[["Name","Score"],["Alice","90"],["Bob","85"]]',
|
|
"spreadsheet_id": "abc123",
|
|
},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-1",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
assert response.success is True
|
|
# Verify the input was coerced from string to list[list[str]]
|
|
assert block._captured_inputs["values"] == [
|
|
["Name", "Score"],
|
|
["Alice", "90"],
|
|
["Bob", "85"],
|
|
]
|
|
assert isinstance(block._captured_inputs["values"], list)
|
|
assert isinstance(block._captured_inputs["values"][0], list)
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_json_string_to_list():
|
|
"""JSON string → list[str]."""
|
|
block = _make_coerce_block(
|
|
"list-block",
|
|
"List Block",
|
|
{"items": list[str]},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="list-block",
|
|
input_data={"items": '["a","b","c"]'},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-2",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
assert block._captured_inputs["items"] == ["a", "b", "c"]
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_json_string_to_dict():
|
|
"""JSON string → dict[str, str]."""
|
|
block = _make_coerce_block(
|
|
"dict-block",
|
|
"Dict Block",
|
|
{"config": dict[str, str]},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="dict-block",
|
|
input_data={"config": '{"key": "value", "foo": "bar"}'},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-3",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
assert block._captured_inputs["config"] == {"key": "value", "foo": "bar"}
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_no_coercion_when_type_matches():
|
|
"""Already-correct types pass through without coercion."""
|
|
block = _make_coerce_block(
|
|
"pass-through",
|
|
"Pass Through",
|
|
{"values": list[list[str]], "name": str},
|
|
)
|
|
|
|
original_values = [["a", "b"], ["c", "d"]]
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="pass-through",
|
|
input_data={"values": original_values, "name": "test"},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-4",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
assert block._captured_inputs["values"] == original_values
|
|
assert block._captured_inputs["name"] == "test"
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_string_to_int():
|
|
"""String number → int."""
|
|
block = _make_coerce_block(
|
|
"int-block",
|
|
"Int Block",
|
|
{"count": int},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="int-block",
|
|
input_data={"count": "42"},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-5",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
assert block._captured_inputs["count"] == 42
|
|
assert isinstance(block._captured_inputs["count"], int)
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_skips_none_values():
|
|
"""None values are not coerced (they may be optional fields)."""
|
|
block = _make_coerce_block(
|
|
"optional-block",
|
|
"Optional Block",
|
|
{"data": list[str], "label": str},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="optional-block",
|
|
input_data={"label": "test"},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-6",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
# 'data' was not provided, so it should not appear in captured inputs
|
|
assert "data" not in block._captured_inputs
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_union_type_preserves_valid_member():
|
|
"""Union-typed fields should not be coerced when the value matches a member."""
|
|
block = _make_coerce_block(
|
|
"union-block",
|
|
"Union Block",
|
|
{"content": str | list[str]},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="union-block",
|
|
input_data={"content": ["a", "b"]},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-7",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
# list[str] should NOT be stringified to '["a", "b"]'
|
|
assert block._captured_inputs["content"] == ["a", "b"]
|
|
assert isinstance(block._captured_inputs["content"], list)
|
|
|
|
|
|
@pytest.mark.asyncio(loop_scope="session")
|
|
async def test_coerce_inner_elements_of_generic():
|
|
"""Inner elements of generic containers are recursively coerced."""
|
|
block = _make_coerce_block(
|
|
"inner-coerce",
|
|
"Inner Coerce",
|
|
{"values": list[str]},
|
|
)
|
|
|
|
mock_workspace_db = MagicMock()
|
|
mock_workspace_db.get_or_create_workspace = AsyncMock(
|
|
return_value=MagicMock(id="ws-1")
|
|
)
|
|
|
|
with patch(
|
|
"backend.copilot.tools.helpers.workspace_db",
|
|
return_value=mock_workspace_db,
|
|
):
|
|
response = await execute_block(
|
|
block=block,
|
|
block_id="inner-coerce",
|
|
# Inner elements are ints, but target is list[str]
|
|
input_data={"values": [1, 2, 3]},
|
|
user_id=_TEST_USER_ID,
|
|
session_id=_TEST_SESSION_ID,
|
|
node_exec_id="exec-8",
|
|
matched_credentials={},
|
|
)
|
|
|
|
assert isinstance(response, BlockOutputResponse)
|
|
# Inner elements should be coerced from int to str
|
|
assert block._captured_inputs["values"] == ["1", "2", "3"]
|
|
assert all(isinstance(v, str) for v in block._captured_inputs["values"])
|