cognee 0.3.4.dev4__py3-none-any.whl → 0.3.5__py3-none-any.whl

cognee/tests/test_pgvector.py

@@ -68,6 +68,44 @@ async def test_getting_of_documents(dataset_name_1):
     )
 
 
+async def test_vector_engine_search_none_limit():
+    file_path_quantum = os.path.join(
+        pathlib.Path(__file__).parent, "test_data/Quantum_computers.txt"
+    )
+
+    file_path_nlp = os.path.join(
+        pathlib.Path(__file__).parent,
+        "test_data/Natural_language_processing.txt",
+    )
+
+    await cognee.prune.prune_data()
+    await cognee.prune.prune_system(metadata=True)
+
+    await cognee.add(file_path_quantum)
+
+    await cognee.add(file_path_nlp)
+
+    await cognee.cognify()
+
+    query_text = "Tell me about Quantum computers"
+
+    from cognee.infrastructure.databases.vector import get_vector_engine
+
+    vector_engine = get_vector_engine()
+
+    collection_name = "Entity_name"
+
+    query_vector = (await vector_engine.embedding_engine.embed_text([query_text]))[0]
+
+    result = await vector_engine.search(
+        collection_name=collection_name, query_vector=query_vector, limit=None
+    )
+
+    # Check that we did not accidentally use any default value for limit
+    # in vector search along the way (like 5, 10, or 15)
+    assert len(result) > 15
+
+
 async def main():
     cognee.config.set_vector_db_config(
         {"vector_db_url": "", "vector_db_key": "", "vector_db_provider": "pgvector"}
@@ -174,6 +212,8 @@ async def main():
     tables_in_database = await vector_engine.get_table_names()
     assert len(tables_in_database) == 0, "PostgreSQL database is not empty"
 
+    await test_vector_engine_search_none_limit()
+
 
 if __name__ == "__main__":
     import asyncio

cognee/tests/test_relational_db_migration.py

@@ -197,6 +197,80 @@ async def relational_db_migration():
     print(f"All checks passed for {graph_db_provider} provider with '{relationship_label}' edges!")
 
 
+async def test_schema_only_migration():
+    # 1. Setup test DB and extract schema
+    migration_engine = await setup_test_db()
+    schema = await migration_engine.extract_schema()
+
+    # 2. Setup graph engine
+    graph_engine = await get_graph_engine()
+
+    # 4. Migrate schema only
+    await migrate_relational_database(graph_engine, schema=schema, schema_only=True)
+
+    # 5. Verify number of tables through search
+    search_results = await cognee.search(
+        query_text="How many tables are there in this database",
+        query_type=cognee.SearchType.GRAPH_COMPLETION,
+        top_k=30,
+    )
+    assert any("11" in r for r in search_results), (
+        "Number of tables in the database reported in search_results is either None or not equal to 11"
+    )
+
+    graph_db_provider = os.getenv("GRAPH_DATABASE_PROVIDER", "networkx").lower()
+
+    edge_counts = {
+        "is_part_of": 0,
+        "has_relationship": 0,
+        "foreign_key": 0,
+    }
+
+    if graph_db_provider == "neo4j":
+        for rel_type in edge_counts.keys():
+            query_str = f"""
+            MATCH ()-[r:{rel_type}]->()
+            RETURN count(r) as c
+            """
+            rows = await graph_engine.query(query_str)
+            edge_counts[rel_type] = rows[0]["c"]
+
+    elif graph_db_provider == "kuzu":
+        for rel_type in edge_counts.keys():
+            query_str = f"""
+            MATCH ()-[r:EDGE]->()
+            WHERE r.relationship_name = '{rel_type}'
+            RETURN count(r) as c
+            """
+            rows = await graph_engine.query(query_str)
+            edge_counts[rel_type] = rows[0][0]
+
+    elif graph_db_provider == "networkx":
+        nodes, edges = await graph_engine.get_graph_data()
+        for _, _, key, _ in edges:
+            if key in edge_counts:
+                edge_counts[key] += 1
+
+    else:
+        raise ValueError(f"Unsupported graph database provider: {graph_db_provider}")
+
+    # 7. Assert counts match expected values
+    expected_counts = {
+        "is_part_of": 11,
+        "has_relationship": 22,
+        "foreign_key": 11,
+    }
+
+    for rel_type, expected in expected_counts.items():
+        actual = edge_counts[rel_type]
+        assert actual == expected, (
+            f"Expected {expected} edges for relationship '{rel_type}', but found {actual}"
+        )
+
+    print("Schema-only migration edge counts validated successfully!")
+    print(f"Edge counts: {edge_counts}")
+
+
 async def test_migration_sqlite():
     database_to_migrate_path = os.path.join(pathlib.Path(__file__).parent, "test_data/")
 
@@ -209,6 +283,7 @@ async def test_migration_sqlite():
     )
 
     await relational_db_migration()
+    await test_schema_only_migration()
 
 
 async def test_migration_postgres():
@@ -224,6 +299,7 @@ async def test_migration_postgres():
         }
     )
     await relational_db_migration()
+    await test_schema_only_migration()
 
 
 async def main():

cognee/tests/unit/infrastructure/databases/test_index_graph_edges.py

@@ -1,5 +1,5 @@
 import pytest
-from unittest.mock import AsyncMock, patch
+from unittest.mock import AsyncMock, patch, MagicMock
 from cognee.tasks.storage.index_graph_edges import index_graph_edges
 
 
@@ -16,6 +16,7 @@ async def test_index_graph_edges_success():
         ],
     )
     mock_vector_engine = AsyncMock()
+    mock_vector_engine.embedding_engine.get_batch_size = MagicMock(return_value=100)
 
     # Patch the globals of the function so that when it does:
     #   vector_engine = get_vector_engine()

cognee/tests/unit/modules/ontology/test_ontology_adapter.py

@@ -1,12 +1,14 @@
 import pytest
 from rdflib import Graph, Namespace, RDF, OWL, RDFS
-from cognee.modules.ontology.rdf_xml.OntologyResolver import OntologyResolver, AttachedOntologyNode
+from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+from cognee.modules.ontology.models import AttachedOntologyNode
+from cognee.modules.ontology.get_default_ontology_resolver import get_default_ontology_resolver
 
 
 def test_ontology_adapter_initialization_success():
-    """Test successful initialization of OntologyAdapter."""
+    """Test successful initialization of RDFLibOntologyResolver from get_default_ontology_resolver."""
 
-    adapter = OntologyResolver()
+    adapter = get_default_ontology_resolver()
     adapter.build_lookup()
 
     assert isinstance(adapter.lookup, dict)
@@ -14,7 +16,7 @@ def test_ontology_adapter_initialization_success():
 
 def test_ontology_adapter_initialization_file_not_found():
     """Test OntologyAdapter initialization with nonexistent file."""
-    adapter = OntologyResolver(ontology_file="nonexistent.owl")
+    adapter = RDFLibOntologyResolver(ontology_file="nonexistent.owl")
     assert adapter.graph is None
 
 
@@ -27,7 +29,7 @@ def test_build_lookup():
 
     g.add((ns.Audi, RDF.type, ns.Car))
 
-    resolver = OntologyResolver()
+    resolver = RDFLibOntologyResolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -50,7 +52,7 @@ def test_find_closest_match_exact():
     g.add((ns.Car, RDF.type, OWL.Class))
     g.add((ns.Audi, RDF.type, ns.Car))
 
-    resolver = OntologyResolver()
+    resolver = RDFLibOntologyResolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -71,7 +73,7 @@ def test_find_closest_match_fuzzy():
     g.add((ns.Audi, RDF.type, ns.Car))
     g.add((ns.BMW, RDF.type, ns.Car))
 
-    resolver = OntologyResolver()
+    resolver = RDFLibOntologyResolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -92,7 +94,7 @@ def test_find_closest_match_no_match():
     g.add((ns.Audi, RDF.type, ns.Car))
     g.add((ns.BMW, RDF.type, ns.Car))
 
-    resolver = OntologyResolver()
+    resolver = RDFLibOntologyResolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -102,10 +104,10 @@ def test_find_closest_match_no_match():
 
 
 def test_get_subgraph_no_match_rdflib():
-    """Test get_subgraph returns empty results for a non-existent node."""
+    """Test get_subgraph returns empty results for a non-existent node using RDFLibOntologyResolver."""
     g = Graph()
 
-    resolver = OntologyResolver()
+    resolver = get_default_ontology_resolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -138,7 +140,7 @@ def test_get_subgraph_success_rdflib():
     g.add((ns.VW, owns, ns.Audi))
     g.add((ns.VW, owns, ns.Porsche))
 
-    resolver = OntologyResolver()
+    resolver = RDFLibOntologyResolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -160,10 +162,10 @@ def test_get_subgraph_success_rdflib():
 
 
 def test_refresh_lookup_rdflib():
-    """Test that refresh_lookup rebuilds the lookup dict into a new object."""
+    """Test that refresh_lookup rebuilds the lookup dict into a new object using RDFLibOntologyResolver."""
     g = Graph()
 
-    resolver = OntologyResolver()
+    resolver = get_default_ontology_resolver()
     resolver.graph = g
     resolver.build_lookup()
 
@@ -172,3 +174,318 @@ def test_refresh_lookup_rdflib():
     resolver.refresh_lookup()
 
     assert resolver.lookup is not original_lookup
+
+
+def test_fuzzy_matching_strategy_exact_match():
+    """Test FuzzyMatchingStrategy finds exact matches."""
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    strategy = FuzzyMatchingStrategy()
+    candidates = ["audi", "bmw", "mercedes"]
+
+    result = strategy.find_match("audi", candidates)
+    assert result == "audi"
+
+
+def test_fuzzy_matching_strategy_fuzzy_match():
+    """Test FuzzyMatchingStrategy finds fuzzy matches."""
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    strategy = FuzzyMatchingStrategy(cutoff=0.6)
+    candidates = ["audi", "bmw", "mercedes"]
+
+    result = strategy.find_match("audii", candidates)
+    assert result == "audi"
+
+
+def test_fuzzy_matching_strategy_no_match():
+    """Test FuzzyMatchingStrategy returns None when no match meets cutoff."""
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    strategy = FuzzyMatchingStrategy(cutoff=0.9)
+    candidates = ["audi", "bmw", "mercedes"]
+
+    result = strategy.find_match("completely_different", candidates)
+    assert result is None
+
+
+def test_fuzzy_matching_strategy_empty_candidates():
+    """Test FuzzyMatchingStrategy handles empty candidates list."""
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    strategy = FuzzyMatchingStrategy()
+
+    result = strategy.find_match("audi", [])
+    assert result is None
+
+
+def test_base_ontology_resolver_initialization():
+    """Test BaseOntologyResolver initialization with default matching strategy."""
+    from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    class TestOntologyResolver(BaseOntologyResolver):
+        def build_lookup(self):
+            pass
+
+        def refresh_lookup(self):
+            pass
+
+        def find_closest_match(self, name, category):
+            return None
+
+        def get_subgraph(self, node_name, node_type="individuals", directed=True):
+            return [], [], None
+
+    resolver = TestOntologyResolver()
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+
+
+def test_base_ontology_resolver_custom_matching_strategy():
+    """Test BaseOntologyResolver initialization with custom matching strategy."""
+    from cognee.modules.ontology.base_ontology_resolver import BaseOntologyResolver
+    from cognee.modules.ontology.matching_strategies import MatchingStrategy
+
+    class CustomMatchingStrategy(MatchingStrategy):
+        def find_match(self, name, candidates):
+            return "custom_match"
+
+    class TestOntologyResolver(BaseOntologyResolver):
+        def build_lookup(self):
+            pass
+
+        def refresh_lookup(self):
+            pass
+
+        def find_closest_match(self, name, category):
+            return None
+
+        def get_subgraph(self, node_name, node_type="individuals", directed=True):
+            return [], [], None
+
+    custom_strategy = CustomMatchingStrategy()
+    resolver = TestOntologyResolver(matching_strategy=custom_strategy)
+    assert resolver.matching_strategy == custom_strategy
+
+
+def test_ontology_config_structure():
+    """Test TypedDict structure for ontology configuration."""
+    from cognee.modules.ontology.ontology_config import Config
+    from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    matching_strategy = FuzzyMatchingStrategy()
+    resolver = RDFLibOntologyResolver(matching_strategy=matching_strategy)
+
+    config: Config = {"ontology_config": {"ontology_resolver": resolver}}
+
+    assert config["ontology_config"]["ontology_resolver"] == resolver
+
+
+def test_get_ontology_resolver_default():
+    """Test get_default_ontology_resolver returns a properly configured RDFLibOntologyResolver with FuzzyMatchingStrategy."""
+    from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    resolver = get_default_ontology_resolver()
+
+    assert isinstance(resolver, RDFLibOntologyResolver)
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+
+
+def test_get_default_ontology_resolver():
+    """Test get_default_ontology_resolver returns a properly configured RDFLibOntologyResolver with FuzzyMatchingStrategy."""
+    from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    resolver = get_default_ontology_resolver()
+
+    assert isinstance(resolver, RDFLibOntologyResolver)
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+
+
+def test_rdflib_ontology_resolver_uses_matching_strategy():
+    """Test that RDFLibOntologyResolver uses the provided matching strategy."""
+    from cognee.modules.ontology.matching_strategies import MatchingStrategy
+
+    class TestMatchingStrategy(MatchingStrategy):
+        def find_match(self, name, candidates):
+            return "test_match" if candidates else None
+
+    ns = Namespace("http://example.org/test#")
+    g = Graph()
+    g.add((ns.Car, RDF.type, OWL.Class))
+    g.add((ns.Audi, RDF.type, ns.Car))
+
+    resolver = RDFLibOntologyResolver(matching_strategy=TestMatchingStrategy())
+    resolver.graph = g
+    resolver.build_lookup()
+
+    result = resolver.find_closest_match("Audi", "individuals")
+    assert result == "test_match"
+
+
+def test_rdflib_ontology_resolver_default_matching_strategy():
+    """Test that RDFLibOntologyResolver uses FuzzyMatchingStrategy by default."""
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    resolver = RDFLibOntologyResolver()
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+
+
+def test_get_ontology_resolver_from_env_success():
+    """Test get_ontology_resolver_from_env returns correct resolver with valid parameters."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+    from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    resolver = get_ontology_resolver_from_env(
+        ontology_resolver="rdflib", matching_strategy="fuzzy", ontology_file_path="/test/path.owl"
+    )
+
+    assert isinstance(resolver, RDFLibOntologyResolver)
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+    assert resolver.ontology_file == "/test/path.owl"
+
+
+def test_get_ontology_resolver_from_env_unsupported_resolver():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for unsupported resolver."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="unsupported",
+            matching_strategy="fuzzy",
+            ontology_file_path="/test/path.owl",
+        )
+
+    assert "Unsupported ontology resolver: unsupported" in str(exc_info.value)
+    assert "Supported resolvers are: RdfLib with FuzzyMatchingStrategy" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_unsupported_strategy():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for unsupported strategy."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="rdflib",
+            matching_strategy="unsupported",
+            ontology_file_path="/test/path.owl",
+        )
+
+    assert "Unsupported ontology resolver: rdflib" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_empty_file_path():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for empty file path."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="rdflib", matching_strategy="fuzzy", ontology_file_path=""
+        )
+
+    assert "Unsupported ontology resolver: rdflib" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_none_file_path():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for None file path."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="rdflib", matching_strategy="fuzzy", ontology_file_path=None
+        )
+
+    assert "Unsupported ontology resolver: rdflib" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_empty_resolver():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for empty resolver."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="", matching_strategy="fuzzy", ontology_file_path="/test/path.owl"
+        )
+
+    assert "Unsupported ontology resolver:" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_empty_strategy():
+    """Test get_ontology_resolver_from_env raises EnvironmentError for empty strategy."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env(
+            ontology_resolver="rdflib", matching_strategy="", ontology_file_path="/test/path.owl"
+        )
+
+    assert "Unsupported ontology resolver: rdflib" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_default_parameters():
+    """Test get_ontology_resolver_from_env with default empty parameters raises EnvironmentError."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError) as exc_info:
+        get_ontology_resolver_from_env()
+
+    assert "Unsupported ontology resolver:" in str(exc_info.value)
+
+
+def test_get_ontology_resolver_from_env_case_sensitivity():
+    """Test get_ontology_resolver_from_env is case sensitive."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    with pytest.raises(EnvironmentError):
+        get_ontology_resolver_from_env(
+            ontology_resolver="RDFLIB",
+            matching_strategy="fuzzy",
+            ontology_file_path="/test/path.owl",
+        )
+
+    with pytest.raises(EnvironmentError):
+        get_ontology_resolver_from_env(
+            ontology_resolver="RdfLib",
+            matching_strategy="fuzzy",
+            ontology_file_path="/test/path.owl",
+        )
+
+
+def test_get_ontology_resolver_from_env_with_actual_file():
+    """Test get_ontology_resolver_from_env works with actual file path."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+    from cognee.modules.ontology.rdf_xml.RDFLibOntologyResolver import RDFLibOntologyResolver
+    from cognee.modules.ontology.matching_strategies import FuzzyMatchingStrategy
+
+    resolver = get_ontology_resolver_from_env(
+        ontology_resolver="rdflib",
+        matching_strategy="fuzzy",
+        ontology_file_path="/path/to/ontology.owl",
+    )
+
+    assert isinstance(resolver, RDFLibOntologyResolver)
+    assert isinstance(resolver.matching_strategy, FuzzyMatchingStrategy)
+    assert resolver.ontology_file == "/path/to/ontology.owl"
+
+
+def test_get_ontology_resolver_from_env_resolver_functionality():
+    """Test that resolver created from env function works correctly."""
+    from cognee.modules.ontology.get_default_ontology_resolver import get_ontology_resolver_from_env
+
+    resolver = get_ontology_resolver_from_env(
+        ontology_resolver="rdflib", matching_strategy="fuzzy", ontology_file_path="/test/path.owl"
+    )
+
+    resolver.build_lookup()
+    assert isinstance(resolver.lookup, dict)
+
+    result = resolver.find_closest_match("test", "individuals")
+    assert result is None  # Should return None for non-existent entity
+
+    nodes, relationships, start_node = resolver.get_subgraph("test", "individuals")
+    assert nodes == []
+    assert relationships == []
+    assert start_node is None

cognee/tests/unit/modules/retrieval/graph_completion_retriever_context_extension_test.py

@@ -7,7 +7,6 @@ import cognee
 from cognee.low_level import setup, DataPoint
 from cognee.tasks.storage import add_data_points
 from cognee.modules.graph.utils import resolve_edges_to_text
-from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_context_extension_retriever import (
     GraphCompletionContextExtensionRetriever,
 )
@@ -165,9 +164,6 @@ class TestGraphCompletionWithContextExtensionRetriever:
 
         retriever = GraphCompletionContextExtensionRetriever()
 
-        with pytest.raises(DatabaseNotCreatedError):
-            await retriever.get_context("Who works at Figma?")
-
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")

cognee/tests/unit/modules/retrieval/graph_completion_retriever_cot_test.py

@@ -7,7 +7,6 @@ import cognee
 from cognee.low_level import setup, DataPoint
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
-from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_cot_retriever import GraphCompletionCotRetriever
 
 
@@ -158,9 +157,6 @@ class TestGraphCompletionCoTRetriever:
 
         retriever = GraphCompletionCotRetriever()
 
-        with pytest.raises(DatabaseNotCreatedError):
-            await retriever.get_context("Who works at Figma?")
-
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")

cognee/tests/unit/modules/retrieval/graph_completion_retriever_test.py

@@ -7,7 +7,6 @@ import cognee
 from cognee.low_level import setup, DataPoint
 from cognee.modules.graph.utils import resolve_edges_to_text
 from cognee.tasks.storage import add_data_points
-from cognee.infrastructure.databases.exceptions import DatabaseNotCreatedError
 from cognee.modules.retrieval.graph_completion_retriever import GraphCompletionRetriever
 
 
@@ -218,9 +217,6 @@ class TestGraphCompletionRetriever:
 
         retriever = GraphCompletionRetriever()
 
-        with pytest.raises(DatabaseNotCreatedError):
-            await retriever.get_context("Who works at Figma?")
-
         await setup()
 
         context = await retriever.get_context("Who works at Figma?")