Double validation for affinity_propagation public function

Ala-Na · Ala-Na · commit 258dff1276e2 · 2022-12-08T22:18:16.000+01:00
diff --git a/sklearn/cluster/_affinity_propagation.py b/sklearn/cluster/_affinity_propagation.py
@@ -13,7 +13,7 @@
 from ..exceptions import ConvergenceWarning
 from ..base import BaseEstimator, ClusterMixin
 from ..utils import as_float_array, check_random_state
-from ..utils._param_validation import Interval, StrOptions
+from ..utils._param_validation import Interval, StrOptions, validate_params
 from ..utils.validation import check_is_fitted
 from ..metrics import euclidean_distances
 from ..metrics import pairwise_distances_argmin
@@ -178,6 +178,12 @@ def _affinity_propagation(
 # Public API
 
 
+@validate_params(
+    {
+        "S": ["array-like"],
+        "return_n_iter": ["boolean"],
+    }
+)
 def affinity_propagation(
     S,
     *,
diff --git a/sklearn/tests/test_public_functions.py b/sklearn/tests/test_public_functions.py
@@ -105,6 +105,7 @@ def _check_function_param_validation(
     "sklearn.metrics.zero_one_loss",
     "sklearn.model_selection.train_test_split",
     "sklearn.svm.l1_min_c",
+    "sklearn.cluster.affinity_propagation",
 ]
 
 
