def test_update_membership():
test_model = KMeans(n_clust=2, random_seed=0)
test_x_train = np.array([[1, 2], [1, 4], [1, 0], [10, 2], [10, 4], [10,
0]])
test_model.classes = np.array([1, 0, 1, 0, 1, 0])
test_model.centers = np.array([[7., 2.], [4., 2.]])
test_model._update_membership(test_x_train)
expected_classes = np.array([1, 1, 1, 0, 0, 0])
np.testing.assert_array_equal(test_model.classes, expected_classes)
def test_predict():
test_model = KMeans(n_clust=2, random_seed=0)
x_test = np.array([[1, 2], [30, 2]])
test_model.centers = np.array([[7., 2.], [4., 2.]])
expected_classes = np.array([1, 0])
np.testing.assert_array_equal(test_model.predict(x_test), expected_classes)
