"""Unit tests for ToolWrapperInputHandler.""" import pytest from langchain_core.tools import StructuredTool from langflow_stepflow.worker.handlers.tool_wrapper import ( ToolWrapperInputHandler, _create_tool_from_wrapper, _execute_calculator_tool, ) def _make_tool_wrapper( *, name: str = "test_tool", description: str = "A test tool", component_code: str | None = "print('hello')", code_blob_id: str | None = None, properties: dict | None = None, static_inputs: dict | None = None, component_type: str = "TestComponent", session_id: str = "test_session", ) -> dict: """Create a tool wrapper dict for testing.""" wrapper: dict = { "__tool_wrapper__": True, "tool_metadata": {"name": name, "description": description}, "tool_input_schema": {"properties": properties or {}}, "static_inputs": static_inputs or {}, "component_type": component_type, "session_id": session_id, } if component_code is not None: wrapper["component_code"] = component_code if code_blob_id is not None: wrapper["code_blob_id"] = code_blob_id return wrapper # --------------------------------------------------------------------------- # _execute_calculator_tool # --------------------------------------------------------------------------- class TestExecuteCalculatorTool: def test_simple_addition(self): assert _execute_calculator_tool("2 + 3") == "5" def test_multiplication(self): assert _execute_calculator_tool("6 * 7") == "42" def test_division(self): assert _execute_calculator_tool("10 / 4") == "2.5" def test_power(self): assert _execute_calculator_tool("2 ** 10") == "1024" def test_complex_expression(self): assert _execute_calculator_tool("(2 + 3) * 4") == "20" def test_float_result(self): assert _execute_calculator_tool("1 / 3") == "0.333333" def test_invalid_expression(self): result = _execute_calculator_tool("not_valid") assert "Calculator error" in result def test_division_by_zero(self): result = _execute_calculator_tool("1 / 0") assert "Calculator error" in result # --------------------------------------------------------------------------- # _create_tool_from_wrapper # --------------------------------------------------------------------------- class TestCreateToolFromWrapper: def test_creates_structured_tool(self): wrapper = _make_tool_wrapper(name="my_tool", description="does stuff") tool = _create_tool_from_wrapper(wrapper) assert isinstance(tool, StructuredTool) assert tool.name == "my_tool" assert tool.description == "does stuff" def test_creates_tool_with_code_blob_id(self): wrapper = _make_tool_wrapper( component_code=None, code_blob_id="abc123", ) tool = _create_tool_from_wrapper(wrapper) assert isinstance(tool, StructuredTool) def test_tool_with_input_schema(self): wrapper = _make_tool_wrapper( properties={ "query": {"type": "string", "default": ""}, "limit": {"type": "integer", "default": "10"}, }, ) tool = _create_tool_from_wrapper(wrapper) assert isinstance(tool, StructuredTool) def test_calculator_tool_execution(self): wrapper = _make_tool_wrapper( name="evaluate_expression", properties={"expression": {"type": "string", "default": ""}}, ) tool = _create_tool_from_wrapper(wrapper) result = tool.invoke({"expression": "2 + 3"}) assert result == {"result": "5"} def test_non_calculator_tool_execution(self): wrapper = _make_tool_wrapper( name="search_tool", properties={"query": {"type": "string", "default": ""}}, static_inputs={"api_key": "test_key"}, # pragma: allowlist secret component_type="SearchComponent", session_id="sess_1", ) tool = _create_tool_from_wrapper(wrapper) result = tool.invoke({"query": "hello"}) assert result["component_type"] == "SearchComponent" assert result["inputs"]["query"] == "hello" assert result["inputs"]["api_key"] == "test_key" # pragma: allowlist secret assert result["inputs"]["session_id"] == "sess_1" assert result["status"] == "tool_wrapper_execution" def test_tool_with_code_blob_id_in_result(self): wrapper = _make_tool_wrapper( name="blob_tool", component_code=None, code_blob_id="blob_abc", ) tool = _create_tool_from_wrapper(wrapper) result = tool.invoke({}) assert result["code_blob_id"] == "blob_abc" def test_tool_with_component_code_in_result(self): wrapper = _make_tool_wrapper( name="code_tool", component_code="def run(): pass", ) tool = _create_tool_from_wrapper(wrapper) result = tool.invoke({}) assert result["has_component_code"] is True def test_missing_both_code_sources_returns_failed_wrapper(self): wrapper = _make_tool_wrapper(component_code=None, code_blob_id=None) tool = _create_tool_from_wrapper(wrapper) # Should be a FailedToolWrapper, not a StructuredTool assert not isinstance(tool, StructuredTool) assert hasattr(tool, "name") assert tool.name == "test_tool" result = tool.invoke({}) assert "error" in result assert "Tool creation failed" in result["error"] # --------------------------------------------------------------------------- # ToolWrapperInputHandler.matches # --------------------------------------------------------------------------- class TestToolWrapperInputHandlerMatches: def setup_method(self): self.handler = ToolWrapperInputHandler() def test_matches_dict_with_tool_wrapper_marker(self): wrapper = _make_tool_wrapper() assert self.handler.matches(template_field={}, value=wrapper) is True def test_matches_list_with_tool_wrapper(self): wrapper = _make_tool_wrapper() assert self.handler.matches(template_field={}, value=[wrapper]) is True def test_matches_list_with_mixed_items(self): wrapper = _make_tool_wrapper() assert self.handler.matches(template_field={}, value=[wrapper, "other"]) is True def test_no_match_plain_dict(self): assert self.handler.matches(template_field={}, value={"key": "value"}) is False def test_no_match_dict_with_false_marker(self): assert self.handler.matches(template_field={}, value={"__tool_wrapper__": False}) is False def test_no_match_string(self): assert self.handler.matches(template_field={}, value="hello") is False def test_no_match_empty_list(self): assert self.handler.matches(template_field={}, value=[]) is False def test_no_match_list_without_wrappers(self): assert self.handler.matches(template_field={}, value=["a", "b", 42]) is False def test_no_match_none(self): assert self.handler.matches(template_field={}, value=None) is False # --------------------------------------------------------------------------- # ToolWrapperInputHandler.prepare # --------------------------------------------------------------------------- class TestToolWrapperInputHandlerPrepare: def setup_method(self): self.handler = ToolWrapperInputHandler() @pytest.mark.asyncio async def test_prepare_single_wrapper(self): wrapper = _make_tool_wrapper(name="my_tool") fields = {"tools": (wrapper, {})} result = await self.handler.prepare(fields, None) assert "tools" in result assert isinstance(result["tools"], StructuredTool) assert result["tools"].name == "my_tool" @pytest.mark.asyncio async def test_prepare_list_of_wrappers(self): wrapper_a = _make_tool_wrapper(name="tool_a") wrapper_b = _make_tool_wrapper(name="tool_b") fields = {"tools": ([wrapper_a, wrapper_b], {})} result = await self.handler.prepare(fields, None) assert "tools" in result assert len(result["tools"]) == 2 assert isinstance(result["tools"][0], StructuredTool) assert isinstance(result["tools"][1], StructuredTool) assert result["tools"][0].name == "tool_a" assert result["tools"][1].name == "tool_b" @pytest.mark.asyncio async def test_prepare_mixed_list(self): wrapper = _make_tool_wrapper(name="real_tool") fields = {"tools": ([wrapper, "not_a_wrapper", 42], {})} result = await self.handler.prepare(fields, None) assert "tools" in result assert len(result["tools"]) == 3 assert isinstance(result["tools"][0], StructuredTool) assert result["tools"][1] == "not_a_wrapper" assert result["tools"][2] == 42 @pytest.mark.asyncio async def test_prepare_multiple_fields(self): wrapper_a = _make_tool_wrapper(name="tool_a") wrapper_b = _make_tool_wrapper(name="tool_b") fields = { "primary_tool": (wrapper_a, {}), "secondary_tool": (wrapper_b, {}), } result = await self.handler.prepare(fields, None) assert isinstance(result["primary_tool"], StructuredTool) assert isinstance(result["secondary_tool"], StructuredTool) assert result["primary_tool"].name == "tool_a" assert result["secondary_tool"].name == "tool_b" @pytest.mark.asyncio async def test_prepare_failed_wrapper_in_list(self): good_wrapper = _make_tool_wrapper(name="good_tool") bad_wrapper = _make_tool_wrapper(name="bad_tool", component_code=None, code_blob_id=None) fields = {"tools": ([good_wrapper, bad_wrapper], {})} result = await self.handler.prepare(fields, None) assert isinstance(result["tools"][0], StructuredTool) # Bad wrapper becomes FailedToolWrapper assert not isinstance(result["tools"][1], StructuredTool) assert result["tools"][1].name == "bad_tool"