"""This file contains tests for custom (user supplied) maskers."""

import numpy as np

import shap


def test_raw_function():
    """Make sure passing a simple masking function works."""
    X, _ = shap.datasets.california(n_points=500)

    def test(X):
        return np.sum(X, 1)

    def custom_masker(mask, x):
        return (x * mask).reshape(1, len(x))  # just zero out the features we are masking

    explainer = shap.Explainer(test, custom_masker)
    shap_values = explainer(X[:100])

    assert np.var(shap_values.values - shap_values.data) < 1e-6  # type: ignore[union-attr, union-attr]