fix: default remote code execution in CSV agent by Adam-Aghili · Pull Request #11762 · langflow-ai/langflow (original) (raw)

400-553: Consider reducing boilerplate across the three security tests with parametrize.

The three test_security_* tests share ~40 lines of identical setup (temp file creation, set_attributes, settings/logger patching, mock agent wiring). A @pytest.mark.parametrize over (allow_dangerous_code_value, expect_warning) — or a small helper that encapsulates the common setup — would make the distinguishing intent of each case clearer and cut repetition.

♻️ Sketch using parametrize

@pytest.mark.parametrize( ("dangerous_code_setting", "expect_warning"), [ (None, False), # default – attribute not set (False, False), # explicitly disabled (True, True), # explicitly enabled ], ids=["default_safe", "explicit_disable", "explicit_enable"], ) def test_security_allow_dangerous_code( self, component_class, mock_langchain_experimental, dangerous_code_setting, expect_warning ): """Verify allow_dangerous_code flag propagation and warning behaviour.""" component = component_class()

with tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False) as f:
    f.write("col1,col2\n1,a\n2,b")
    csv_file = f.name

try:
    attrs = {
        "llm": MagicMock(),
        "path": csv_file,
        "agent_type": "openai-tools",
        "input_value": "test",
        "verbose": False,
        "handle_parsing_errors": True,
        "pandas_kwargs": {},
    }
    if dangerous_code_setting is not None:
        attrs["allow_dangerous_code"] = dangerous_code_setting
    component.set_attributes(attrs)

    with (
        patch("lfx.components.langchain_utilities.csv_agent.get_settings_service") as mock_get_settings,
        patch("lfx.components.langchain_utilities.csv_agent.logger") as mock_logger,
    ):
        mock_create_agent = mock_langchain_experimental
        mock_settings = MagicMock()
        mock_settings.settings.storage_type = "local"
        mock_get_settings.return_value = mock_settings

        mock_agent = MagicMock()
        mock_agent.invoke.return_value = {"output": "result"}
        mock_create_agent.return_value = mock_agent

        result = component.build_agent_response()

        call_kwargs = mock_create_agent.call_args[1]
        expected_flag = dangerous_code_setting if dangerous_code_setting is not None else False
        assert call_kwargs["allow_dangerous_code"] is expected_flag

        if expect_warning:
            mock_logger.warning.assert_called_once()
            warning_msg = mock_logger.warning.call_args[0][0]
            assert "allow_dangerous_code=True" in warning_msg
        else:
            mock_logger.warning.assert_not_called()

        assert isinstance(result, Message)
finally:
    Path(csv_file).unlink()