1  from typing import Callable
 2  
 3  import pytest
 4  
 5  from evidently.guardrails import GuardException
 6  from evidently.guardrails import PythonFunction
 7  from evidently.guardrails import guard
 8  
 9  
10  def stub_guard_factory(result: bool) -> Callable[[str], bool]:
11      def stub_guard(data: str) -> bool:
12          return result
13  
14      return stub_guard
15  
16  
17  @guard(PythonFunction(stub_guard_factory(True)))
18  def success_operation(input: str) -> str:
19      return input
20  
21  
22  @guard(
23      [
24          PythonFunction(stub_guard_factory(True)),
25      ]
26  )
27  def success_operation_multiple_guards(input: str) -> str:
28      return input
29  
30  
31  @pytest.mark.parametrize(
32      "func,args",
33      [
34          (success_operation, ["test"]),
35          (success_operation_multiple_guards, ["test"]),
36      ],
37  )
38  def test_guard_decorator_success(func, args):
39      assert func(*args) == args[0]
40  
41  
42  @guard(PythonFunction(stub_guard_factory(False)))
43  def failed_operation(input: str) -> str:
44      return input
45  
46  
47  @guard(
48      [
49          PythonFunction(stub_guard_factory(False)),
50          PythonFunction(stub_guard_factory(False)),
51      ]
52  )
53  def failed_operation_multiple_guards(input: str) -> str:
54      return input
55  
56  
57  @pytest.mark.parametrize(
58      "func,args",
59      [
60          (failed_operation, ["test"]),
61          (failed_operation_multiple_guards, ["test"]),
62      ],
63  )
64  def test_guard_decorator_failed(func, args):
65      with pytest.raises(GuardException):
66          func(*args)