1  import pandas as pd
 2  
 3  from evidently.legacy.features.custom_feature import CustomSingleColumnFeature
 4  from evidently.legacy.pipeline.column_mapping import ColumnMapping
 5  from evidently.legacy.utils.data_preprocessing import create_data_definition
 6  
 7  
 8  def test_custom_feature():
 9      def add_two(data: pd.Series) -> pd.Series:
10          return data + 2
11  
12      feature_generator = CustomSingleColumnFeature(column_name="column_1", display_name="cl", func=add_two, name="cf")
13      data = pd.DataFrame(dict(column_1=[1, 2, 3]))
14      result = feature_generator.generate_feature(
15          data=data,
16          data_definition=create_data_definition(None, data, ColumnMapping()),
17      )
18  
19      pd.testing.assert_frame_equal(result, pd.DataFrame(dict(cf=[3, 4, 5])))
20      assert feature_generator.as_column().display_name == "cl"