1  import json
 2  
 3  import pandas as pd
 4  import pytest
 5  from sklearn.metrics import r2_score
 6  
 7  from evidently.legacy.base_metric import InputData
 8  from evidently.legacy.test_suite import TestSuite
 9  from evidently.legacy.tests.custom_test import CustomValueTest
10  
11  
12  def r2_func(data: InputData):
13      return r2_score(data.current_data["target"], data.current_data["prediction"])
14  
15  
16  def test_custom_test() -> None:
17      test_current_dataset = pd.DataFrame(
18          {
19              "target": [0, 1],
20              "prediction": [0, 1],
21          }
22      )
23  
24      # happy path
25      suite = TestSuite(tests=[CustomValueTest(r2_func, eq=1)])
26      suite.run(current_data=test_current_dataset, reference_data=test_current_dataset)
27      suite._inner_suite.raise_for_error()
28      assert suite
29      assert suite.show()
30      assert suite.json()
31  
32      with pytest.raises(ValueError):  # no conditions specified
33          CustomValueTest(func=r2_func, title="R2 score test")
34  
35      # failed condition
36      suite = TestSuite(tests=[CustomValueTest(r2_func, lt=0)])
37      suite.run(current_data=test_current_dataset, reference_data=test_current_dataset)
38      suite._inner_suite.raise_for_error()
39      assert not suite
40  
41  
42  def test_custom_test_json_render() -> None:
43      test_current_dataset = pd.DataFrame(
44          {
45              "target": [0, 1],
46              "prediction": [0, 1],
47          }
48      )
49      suite = TestSuite(tests=[CustomValueTest(r2_func, title="r2_score", eq=1)])
50      suite.run(current_data=test_current_dataset, reference_data=test_current_dataset)
51  
52      result_from_json = json.loads(suite.json())
53      assert result_from_json["summary"]["all_passed"] is True
54      test_info = result_from_json["tests"][0]
55      assert test_info == {
56          "name": "Custom Value test",
57          "description": "Custom function 'r2_score' value is 1.0. The test threshold is eq=1",
58          "status": "SUCCESS",
59          "group": "custom",
60          "parameters": {"condition": {"eq": 1}, "value": 1.0},
61      }