|
|
|
@ -12,6 +12,7 @@ class FakeChain(Chain, BaseModel):
|
|
|
|
|
|
|
|
|
|
be_correct: bool = True
|
|
|
|
|
the_input_keys: List[str] = ["foo"]
|
|
|
|
|
the_output_keys: List[str] = ["bar"]
|
|
|
|
|
|
|
|
|
|
@property
|
|
|
|
|
def input_keys(self) -> List[str]:
|
|
|
|
@ -21,7 +22,7 @@ class FakeChain(Chain, BaseModel):
|
|
|
|
|
@property
|
|
|
|
|
def output_keys(self) -> List[str]:
|
|
|
|
|
"""Output key of bar."""
|
|
|
|
|
return ["bar"]
|
|
|
|
|
return self.the_output_keys
|
|
|
|
|
|
|
|
|
|
def _call(self, inputs: Dict[str, str]) -> Dict[str, str]:
|
|
|
|
|
if self.be_correct:
|
|
|
|
@ -63,3 +64,45 @@ def test_single_input_error() -> None:
|
|
|
|
|
chain = FakeChain(the_input_keys=["foo", "bar"])
|
|
|
|
|
with pytest.raises(ValueError):
|
|
|
|
|
chain("bar")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_single_arg() -> None:
|
|
|
|
|
"""Test run method with single arg."""
|
|
|
|
|
chain = FakeChain()
|
|
|
|
|
output = chain.run("bar")
|
|
|
|
|
assert output == "baz"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_multiple_args_error() -> None:
|
|
|
|
|
"""Test run method with multiple args errors as expected."""
|
|
|
|
|
chain = FakeChain()
|
|
|
|
|
with pytest.raises(ValueError):
|
|
|
|
|
chain.run("bar", "foo")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_kwargs() -> None:
|
|
|
|
|
"""Test run method with kwargs."""
|
|
|
|
|
chain = FakeChain(the_input_keys=["foo", "bar"])
|
|
|
|
|
output = chain.run(foo="bar", bar="foo")
|
|
|
|
|
assert output == "baz"
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_kwargs_error() -> None:
|
|
|
|
|
"""Test run method with kwargs errors as expected."""
|
|
|
|
|
chain = FakeChain(the_input_keys=["foo", "bar"])
|
|
|
|
|
with pytest.raises(ValueError):
|
|
|
|
|
chain.run(foo="bar", baz="foo")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_run_args_and_kwargs_error() -> None:
|
|
|
|
|
"""Test run method with args and kwargs."""
|
|
|
|
|
chain = FakeChain(the_input_keys=["foo", "bar"])
|
|
|
|
|
with pytest.raises(ValueError):
|
|
|
|
|
chain.run("bar", foo="bar")
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test_multiple_output_keys_error() -> None:
|
|
|
|
|
"""Test run with multiple output keys errors as expected."""
|
|
|
|
|
chain = FakeChain(the_output_keys=["foo", "bar"])
|
|
|
|
|
with pytest.raises(ValueError):
|
|
|
|
|
chain.run("bar")
|
|
|
|
|