import pytest
import numpy as np
from privacy_tuner.core.utility.query_error import QueryError

def test_query_error_mean():
    orig = np.array([[1, 2], [3, 4], [5, 6]])
    pert = orig + 0.1
    evaluator = QueryError(query_type='mean')
    utility = evaluator.evaluate(orig, pert)
    assert 0.0 <= utility <= 1.0
    assert utility > 0.9

def test_query_error_sum():
    orig = np.array([[1, 2], [3, 4], [5, 6]])
    pert = orig * 1.1
    evaluator = QueryError(query_type='sum')
    utility = evaluator.evaluate(orig, pert)
    assert 0.0 <= utility <= 1.0
    # 使用 pytest.approx 处理浮点误差
    assert utility == pytest.approx(0.9)

def test_query_error_specific_column():
    orig = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    pert = orig.copy()
    pert[:, 1] = pert[:, 1] + 10
    evaluator = QueryError(columns=1, query_type='mean')
    utility = evaluator.evaluate(orig, pert)
    assert utility < 0.5

def test_query_error_with_task():
    orig = np.array([[1, 2], [3, 4]])
    pert = orig + 0.1
    evaluator = QueryError()
    utility = evaluator.evaluate(orig, pert, task={'query_type': 'sum'})
    assert 0.0 <= utility <= 1.0