test META model, not nested estimator

benedekrozemberczki · Aug 11, 2023 · c7ceb75 · c7ceb75
1 parent 3b7052e
commit c7ceb75
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/test/meta_embedding_test.py b/test/meta_embedding_test.py
@@ -26,9 +26,10 @@ def test_neu():
 
     assert embedding.shape[0] == graph.number_of_nodes()
     assert type(embedding) == np.ndarray
-    assert np.array_equal(model.fit_transform(graph).shape, embedding.shape)
+    assert np.array_equal(meta_model.fit_transform(graph, model).shape,
+                          embedding.shape)
 
     y = embedding[:, 0]
-    z, y_hat = model.fit_transform(graph, y)
+    z, y_hat = meta_model.fit_transform(graph, model, y)
     assert y_hat.shape[0] == graph.number_of_nodes()
     assert np.array_equal(y_hat, y)