diff --git a/sklearn/cluster/tests/test_k_means.py b/sklearn/cluster/tests/test_k_means.py index 4d1e517742b3e..9d283a597ef6d 100644 --- a/sklearn/cluster/tests/test_k_means.py +++ b/sklearn/cluster/tests/test_k_means.py @@ -328,7 +328,7 @@ def test_k_means_fortran_aligned_data(): (4, 300, 1e-1), # loose convergence ]) def test_k_means_fit_predict(algo, dtype, constructor, seed, max_iter, tol): - # check that fit.predict gives same result as fit_predict + # check that predict gives same result as fit_predict or labels_ # There's a very small chance of failure with elkan on unstructured dataset # because predict method uses fast euclidean distances computation which # may cause small numerical instabilities. @@ -342,10 +342,12 @@ def test_k_means_fit_predict(algo, dtype, constructor, seed, max_iter, tol): kmeans = KMeans(algorithm=algo, n_clusters=10, random_state=seed, tol=tol, max_iter=max_iter, n_jobs=1) - labels_1 = kmeans.fit(X).predict(X) - labels_2 = kmeans.fit_predict(X) + labels_1 = kmeans.fit_predict(X) + labels_2 = kmeans.predict(X) + labels_3 = kmeans.labels_ assert_array_equal(labels_1, labels_2) + assert_array_equal(labels_1, labels_3) def test_mb_kmeans_verbose():