Update extractor tests

cbornet · cbornet · commit 7fc176c3f264 · 2024-07-22T16:10:43.000+02:00
diff --git a/libs/langchain/tests/unit_tests/test_gliner_link_extractor.py b/libs/langchain/tests/unit_tests/test_gliner_link_extractor.py
@@ -1,5 +1,5 @@
+from langchain_core.graph_vectorstores import Link
 from ragstack_langchain.graph_store.extractors import GLiNERLinkExtractor
-from ragstack_langchain.graph_store.links import Link
 
 PAGE_1 = """
 Cristiano Ronaldo dos Santos Aveiro (Portuguese pronunciation: [kɾiʃ'tjɐnu
diff --git a/libs/langchain/tests/unit_tests/test_hierarchy_link_extractor.py b/libs/langchain/tests/unit_tests/test_hierarchy_link_extractor.py
@@ -1,5 +1,5 @@
+from langchain_core.graph_vectorstores import Link
 from ragstack_langchain.graph_store.extractors import HierarchyLinkExtractor
-from ragstack_langchain.graph_store.links import Link
 
 PATH_1 = ["Root", "H1", "h2"]
 
diff --git a/libs/langchain/tests/unit_tests/test_html_link_extractor.py b/libs/langchain/tests/unit_tests/test_html_link_extractor.py
@@ -1,6 +1,6 @@
 from bs4 import BeautifulSoup
+from langchain_core.graph_vectorstores import Link
 from ragstack_langchain.graph_store.extractors import HtmlInput, HtmlLinkExtractor
-from ragstack_langchain.graph_store.links import Link
 
 PAGE_1 = """
 <html>
diff --git a/libs/langchain/tests/unit_tests/test_keybert_link_extractor.py b/libs/langchain/tests/unit_tests/test_keybert_link_extractor.py
@@ -1,5 +1,5 @@
+from langchain_core.graph_vectorstores import Link
 from ragstack_langchain.graph_store.extractors import KeybertLinkExtractor
-from ragstack_langchain.graph_store.links import Link
 
 PAGE_1 = """
 Supervised learning is the machine learning task of learning a function that
diff --git a/libs/langchain/tests/unit_tests/test_link_extractor_transformer.py b/libs/langchain/tests/unit_tests/test_link_extractor_transformer.py
@@ -1,4 +1,5 @@
 from langchain_core.documents import Document
+from langchain_core.graph_vectorstores.links import Link, get_links
 from ragstack_langchain.graph_store.extractors import (
     HtmlLinkExtractor,
     LinkExtractorTransformer,
@@ -9,7 +10,6 @@
 from ragstack_langchain.graph_store.extractors.keybert_link_extractor import (
     KeybertLinkExtractor,
 )
-from ragstack_langchain.graph_store.links import Link, get_links
 
 from . import (
     test_gliner_link_extractor,
@@ -40,15 +40,15 @@ def test_html_extractor():
     assert results[0] == doc1
     assert results[1] == doc2
 
-    assert get_links(doc1) == {
+    assert set(get_links(doc1)) == {
         Link.incoming(kind="hyperlink", tag="https://foo.com/bar/"),
         Link.outgoing(kind="hyperlink", tag="https://foo.com/bar/relative"),
         Link.outgoing(kind="hyperlink", tag="https://foo.com/relative-base"),
         Link.outgoing(kind="hyperlink", tag="http://cnn.com"),
         Link.outgoing(kind="hyperlink", tag="https://same.foo"),
     }
 
-    assert get_links(doc2) == {
+    assert set(get_links(doc2)) == {
         Link.incoming(kind="hyperlink", tag="https://foo.com/baz/"),
         Link.outgoing(kind="hyperlink", tag="https://foo.com/bar/"),
     }
@@ -71,15 +71,15 @@ def test_multiple_extractors():
     assert results[0] == doc1
     assert results[1] == doc2
 
-    assert get_links(doc1) == {
+    assert set(get_links(doc1)) == {
         Link(kind="kw", direction="bidir", tag="labeled"),
         Link(kind="kw", direction="bidir", tag="learning"),
         Link(kind="kw", direction="bidir", tag="training"),
         Link(kind="kw", direction="bidir", tag="supervised"),
         Link(kind="kw", direction="bidir", tag="labels"),
     }
 
-    assert get_links(doc2) == {
+    assert set(get_links(doc2)) == {
         Link(kind="kw", direction="bidir", tag="cristiano"),
         Link(kind="kw", direction="bidir", tag="goalscorer"),
         Link(kind="kw", direction="bidir", tag="footballer"),
