import pytest
import numpy as np
from privacy_tuner.core.mechanisms import LaplaceMechanism, GaussianMechanism

def test_laplace_mechanism_shape():
    """测试拉普拉斯机制输出形状与输入一致"""
    mechanism = LaplaceMechanism()
    data = np.random.randn(10, 5)
    params = {'epsilon': 1.0, 'sensitivity': 1.0}
    noisy = mechanism.apply(data, params)
    assert noisy.shape == data.shape
    assert not np.array_equal(noisy, data)  # 应该不同

def test_laplace_mechanism_epsilon_zero():
    """测试 epsilon 非常大时，噪声应趋近于0"""
    mechanism = LaplaceMechanism()
    data = np.random.randn(10, 5)
    params = {'epsilon': 1e6, 'sensitivity': 1.0}
    noisy = mechanism.apply(data, params)
    # 噪声很小，近似相等
    np.testing.assert_almost_equal(noisy, data, decimal=3)

def test_gaussian_mechanism_shape():
    mechanism = GaussianMechanism()
    data = np.random.randn(10, 5)
    params = {'epsilon': 1.0, 'delta': 1e-5, 'sensitivity': 1.0}
    noisy = mechanism.apply(data, params)
    assert noisy.shape == data.shape
    assert not np.array_equal(noisy, data)

def test_gaussian_mechanism_delta_zero():
    """测试 delta 非常小时，噪声较大"""
    mechanism = GaussianMechanism()
    data = np.random.randn(10, 5)
    params = {'epsilon': 1.0, 'delta': 1e-10, 'sensitivity': 1.0}
    noisy = mechanism.apply(data, params)
    # 仅检查是否有输出，不做数值验证
    assert noisy is not None

def test_missing_params():
    mechanism = LaplaceMechanism()
    data = np.random.randn(10, 5)
    with pytest.raises(ValueError):
        mechanism.apply(data, {})  # 缺少 epsilon