"""Test python chain."""

from langchain.chains.python import PythonChain


def test_functionality() -> None:
    """Test correct functionality."""
    chain = PythonChain(input_key="code1", output_key="output1")
    code = "print(1 + 1)"
    output = chain({"code1": code})
    assert output == {"code1": code, "output1": "2\n"}

    # Test with the more user-friendly interface.
    simple_output = chain.run(code)
    assert simple_output == "2\n"