1  import pandas as pd
 2  
 3  from evidently.legacy.features.is_valid_sql_feature import IsValidSQL
 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_is_valid_sql_feature():
 9      feature_generator = IsValidSQL("column_1")
10      data = pd.DataFrame(
11          dict(
12              column_1=[
13                  "SELECT * FROM users",  # Valid SQL (simple query)
14                  "SELECT id, address FROM users; SELECT count(id) FROM users",  # Valid SQL (multiple SQL queries)
15                  "INSERT INTO table",  # Invalid SQL (incomplete query)
16                  "SLECT * FROM users",  # Invalid SQL (typo)
17                  "SLECT * FROM users; SELECT id, address FROM users",  # Invalid SQL (1 invalid sub-query)
18              ]
19          )
20      )
21  
22      result = feature_generator.generate_feature(
23          data=data,
24          data_definition=create_data_definition(None, data, ColumnMapping()),
25      )
26  
27      expected_result = pd.DataFrame(dict(column_1=[True, True, False, False, False]))
28  
29      assert result.equals(expected_result)