cognee 0.5.0__py3-none-any.whl → 0.5.0.dev0__py3-none-any.whl

cognee/tests/unit/api/test_ontology_endpoint.py

@@ -1,28 +1,17 @@
 import pytest
 import uuid
 from fastapi.testclient import TestClient
-from unittest.mock import Mock
+from unittest.mock import patch, Mock, AsyncMock
 from types import SimpleNamespace
+import importlib
 from cognee.api.client import app
-from cognee.modules.users.methods import get_authenticated_user
 
-
-@pytest.fixture(scope="session")
-def test_client():
-    # Keep a single TestClient (and event loop) for the whole module.
-    # Re-creating TestClient repeatedly can break async DB connections (asyncpg loop mismatch).
-    with TestClient(app) as c:
-        yield c
+gau_mod = importlib.import_module("cognee.modules.users.methods.get_authenticated_user")
 
 
 @pytest.fixture
-def client(test_client, mock_default_user):
-    async def override_get_authenticated_user():
-        return mock_default_user
-
-    app.dependency_overrides[get_authenticated_user] = override_get_authenticated_user
-    yield test_client
-    app.dependency_overrides.pop(get_authenticated_user, None)
+def client():
+    return TestClient(app)
 
 
 @pytest.fixture
@@ -36,15 +25,16 @@ def mock_user():
 def mock_default_user():
     """Mock default user for testing."""
     return SimpleNamespace(
-        id=str(uuid.uuid4()),
-        email="default@example.com",
-        is_active=True,
-        tenant_id=str(uuid.uuid4()),
+        id=uuid.uuid4(), email="default@example.com", is_active=True, tenant_id=uuid.uuid4()
     )
 
 
-def test_upload_ontology_success(client):
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_ontology_success(mock_get_default_user, client, mock_default_user):
     """Test successful ontology upload"""
+    import json
+
+    mock_get_default_user.return_value = mock_default_user
     ontology_content = (
         b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
     )
@@ -53,7 +43,7 @@ def test_upload_ontology_success(client):
     response = client.post(
         "/api/v1/ontologies",
         files=[("ontology_file", ("test.owl", ontology_content, "application/xml"))],
-        data={"ontology_key": unique_key, "description": "Test"},
+        data={"ontology_key": json.dumps([unique_key]), "description": json.dumps(["Test"])},
     )
 
     assert response.status_code == 200
@@ -62,8 +52,10 @@ def test_upload_ontology_success(client):
     assert "uploaded_at" in data["uploaded_ontologies"][0]
 
 
-def test_upload_ontology_invalid_file(client):
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_ontology_invalid_file(mock_get_default_user, client, mock_default_user):
     """Test 400 response for non-.owl files"""
+    mock_get_default_user.return_value = mock_default_user
     unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
     response = client.post(
         "/api/v1/ontologies",
@@ -73,10 +65,14 @@ def test_upload_ontology_invalid_file(client):
     assert response.status_code == 400
 
 
-def test_upload_ontology_missing_data(client):
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_ontology_missing_data(mock_get_default_user, client, mock_default_user):
     """Test 400 response for missing file or key"""
+    import json
+
+    mock_get_default_user.return_value = mock_default_user
     # Missing file
-    response = client.post("/api/v1/ontologies", data={"ontology_key": "test"})
+    response = client.post("/api/v1/ontologies", data={"ontology_key": json.dumps(["test"])})
     assert response.status_code == 400
 
     # Missing key
@@ -86,25 +82,33 @@ def test_upload_ontology_missing_data(client):
     assert response.status_code == 400
 
 
-def test_upload_ontology_without_auth_header(client):
-    """Test behavior when no explicit authentication header is provided."""
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_ontology_unauthorized(mock_get_default_user, client, mock_default_user):
+    """Test behavior when default user is provided (no explicit authentication)"""
+    import json
+
     unique_key = f"test_ontology_{uuid.uuid4().hex[:8]}"
+    mock_get_default_user.return_value = mock_default_user
     response = client.post(
         "/api/v1/ontologies",
         files=[("ontology_file", ("test.owl", b"<rdf></rdf>", "application/xml"))],
-        data={"ontology_key": unique_key},
+        data={"ontology_key": json.dumps([unique_key])},
     )
 
+    # The current system provides a default user when no explicit authentication is given
+    # This test verifies the system works with conditional authentication
     assert response.status_code == 200
     data = response.json()
     assert data["uploaded_ontologies"][0]["ontology_key"] == unique_key
     assert "uploaded_at" in data["uploaded_ontologies"][0]
 
 
-def test_upload_multiple_ontologies_in_single_request_is_rejected(client):
-    """Uploading multiple ontology files in a single request should fail."""
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_multiple_ontologies(mock_get_default_user, client, mock_default_user):
+    """Test uploading multiple ontology files in single request"""
     import io
 
+    # Create mock files
     file1_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
     file2_content = b"<rdf:RDF xmlns:rdf='http://www.w3.org/1999/02/22-rdf-syntax-ns#'></rdf:RDF>"
 
@@ -112,16 +116,24 @@ def test_upload_multiple_ontologies_in_single_request_is_rejected(client):
         ("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
         ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
     ]
-    data = {"ontology_key": "vehicles", "description": "Base vehicles"}
+    data = {
+        "ontology_key": '["vehicles", "manufacturers"]',
+        "descriptions": '["Base vehicles", "Car manufacturers"]',
+    }
 
     response = client.post("/api/v1/ontologies", files=files, data=data)
 
-    assert response.status_code == 400
-    assert "Only one ontology_file is allowed" in response.json()["error"]
+    assert response.status_code == 200
+    result = response.json()
+    assert "uploaded_ontologies" in result
+    assert len(result["uploaded_ontologies"]) == 2
+    assert result["uploaded_ontologies"][0]["ontology_key"] == "vehicles"
+    assert result["uploaded_ontologies"][1]["ontology_key"] == "manufacturers"
 
 
-def test_upload_endpoint_rejects_array_style_fields(client):
-    """Array-style form values should be rejected (no backwards compatibility)."""
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_upload_endpoint_accepts_arrays(mock_get_default_user, client, mock_default_user):
+    """Test that upload endpoint accepts array parameters"""
     import io
     import json
 
@@ -130,16 +142,18 @@ def test_upload_endpoint_rejects_array_style_fields(client):
     files = [("ontology_file", ("single.owl", io.BytesIO(file_content), "application/xml"))]
     data = {
         "ontology_key": json.dumps(["single_key"]),
-        "description": json.dumps(["Single ontology"]),
+        "descriptions": json.dumps(["Single ontology"]),
     }
 
     response = client.post("/api/v1/ontologies", files=files, data=data)
 
-    assert response.status_code == 400
-    assert "ontology_key must be a string" in response.json()["error"]
+    assert response.status_code == 200
+    result = response.json()
+    assert result["uploaded_ontologies"][0]["ontology_key"] == "single_key"
 
 
-def test_cognify_with_multiple_ontologies(client):
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_cognify_with_multiple_ontologies(mock_get_default_user, client, mock_default_user):
     """Test cognify endpoint accepts multiple ontology keys"""
     payload = {
         "datasets": ["test_dataset"],
@@ -153,11 +167,13 @@ def test_cognify_with_multiple_ontologies(client):
     assert response.status_code in [200, 400, 409]  # May fail for other reasons, not type
 
 
-def test_complete_multifile_workflow(client):
-    """Test workflow: upload ontologies one-by-one → cognify with multiple keys"""
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_complete_multifile_workflow(mock_get_default_user, client, mock_default_user):
+    """Test complete workflow: upload multiple ontologies → cognify with multiple keys"""
     import io
+    import json
 
-    # Step 1: Upload two ontologies (one-by-one)
+    # Step 1: Upload multiple ontologies
     file1_content = b"""<?xml version="1.0"?>
     <rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#"
              xmlns:owl="http://www.w3.org/2002/07/owl#">
@@ -170,21 +186,17 @@ def test_complete_multifile_workflow(client):
         <owl:Class rdf:ID="Manufacturer"/>
     </rdf:RDF>"""
 
-    upload_response_1 = client.post(
-        "/api/v1/ontologies",
-        files=[("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml"))],
-        data={"ontology_key": "vehicles", "description": "Vehicle ontology"},
-    )
-    assert upload_response_1.status_code == 200
+    files = [
+        ("ontology_file", ("vehicles.owl", io.BytesIO(file1_content), "application/xml")),
+        ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml")),
+    ]
+    data = {
+        "ontology_key": json.dumps(["vehicles", "manufacturers"]),
+        "descriptions": json.dumps(["Vehicle ontology", "Manufacturer ontology"]),
+    }
 
-    upload_response_2 = client.post(
-        "/api/v1/ontologies",
-        files=[
-            ("ontology_file", ("manufacturers.owl", io.BytesIO(file2_content), "application/xml"))
-        ],
-        data={"ontology_key": "manufacturers", "description": "Manufacturer ontology"},
-    )
-    assert upload_response_2.status_code == 200
+    upload_response = client.post("/api/v1/ontologies", files=files, data=data)
+    assert upload_response.status_code == 200
 
     # Step 2: Verify ontologies are listed
     list_response = client.get("/api/v1/ontologies")
@@ -205,41 +217,41 @@ def test_complete_multifile_workflow(client):
     assert cognify_response.status_code != 400  # Not a validation error
 
 
-def test_upload_error_handling(client):
-    """Test error handling for invalid uploads (single-file endpoint)."""
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_multifile_error_handling(mock_get_default_user, client, mock_default_user):
+    """Test error handling for invalid multifile uploads"""
     import io
     import json
 
-    # Array-style key should be rejected
+    # Test mismatched array lengths
     file_content = b"<rdf:RDF></rdf:RDF>"
     files = [("ontology_file", ("test.owl", io.BytesIO(file_content), "application/xml"))]
     data = {
-        "ontology_key": json.dumps(["key1", "key2"]),
-        "description": "desc1",
+        "ontology_key": json.dumps(["key1", "key2"]),  # 2 keys, 1 file
+        "descriptions": json.dumps(["desc1"]),
     }
 
     response = client.post("/api/v1/ontologies", files=files, data=data)
     assert response.status_code == 400
-    assert "ontology_key must be a string" in response.json()["error"]
+    assert "Number of keys must match number of files" in response.json()["error"]
 
-    # Duplicate key should be rejected
-    response_1 = client.post(
-        "/api/v1/ontologies",
-        files=[("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml"))],
-        data={"ontology_key": "duplicate", "description": "desc1"},
-    )
-    assert response_1.status_code == 200
+    # Test duplicate keys
+    files = [
+        ("ontology_file", ("test1.owl", io.BytesIO(file_content), "application/xml")),
+        ("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml")),
+    ]
+    data = {
+        "ontology_key": json.dumps(["duplicate", "duplicate"]),
+        "descriptions": json.dumps(["desc1", "desc2"]),
+    }
 
-    response_2 = client.post(
-        "/api/v1/ontologies",
-        files=[("ontology_file", ("test2.owl", io.BytesIO(file_content), "application/xml"))],
-        data={"ontology_key": "duplicate", "description": "desc2"},
-    )
-    assert response_2.status_code == 400
-    assert "already exists" in response_2.json()["error"]
+    response = client.post("/api/v1/ontologies", files=files, data=data)
+    assert response.status_code == 400
+    assert "Duplicate ontology keys not allowed" in response.json()["error"]
 
 
-def test_cognify_missing_ontology_key(client):
+@patch.object(gau_mod, "get_default_user", new_callable=AsyncMock)
+def test_cognify_missing_ontology_key(mock_get_default_user, client, mock_default_user):
     """Test cognify with non-existent ontology key"""
     payload = {
         "datasets": ["test_dataset"],

cognee/tests/unit/infrastructure/mock_embedding_engine.py

@@ -4,7 +4,10 @@ from typing import List
 from cognee.infrastructure.databases.vector.embeddings.LiteLLMEmbeddingEngine import (
     LiteLLMEmbeddingEngine,
 )
-from cognee.shared.rate_limiting import embedding_rate_limiter_context_manager
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
+    embedding_rate_limit_async,
+    embedding_sleep_and_retry_async,
+)
 
 
 class MockEmbeddingEngine(LiteLLMEmbeddingEngine):
@@ -31,6 +34,8 @@ class MockEmbeddingEngine(LiteLLMEmbeddingEngine):
         self.fail_every_n_requests = fail_every_n_requests
         self.add_delay = add_delay
 
+    @embedding_sleep_and_retry_async()
+    @embedding_rate_limit_async
     async def embed_text(self, text: List[str]) -> List[List[float]]:
         """
         Mock implementation that returns fixed embeddings and can
@@ -47,5 +52,4 @@ class MockEmbeddingEngine(LiteLLMEmbeddingEngine):
             raise Exception(f"Mock failure on request #{self.request_count}")
 
         # Return mock embeddings of the correct dimension
-        async with embedding_rate_limiter_context_manager():
-            return [[0.1] * self.dimensions for _ in text]
+        return [[0.1] * self.dimensions for _ in text]

cognee/tests/unit/infrastructure/test_embedding_rate_limiting_realistic.py

@@ -6,6 +6,9 @@ import logging
 from cognee.infrastructure.llm.config import (
     get_llm_config,
 )
+from cognee.infrastructure.databases.vector.embeddings.embedding_rate_limiter import (
+    EmbeddingRateLimiter,
+)
 from cognee.tests.unit.infrastructure.mock_embedding_engine import MockEmbeddingEngine
 
 # Configure logging
@@ -30,6 +33,7 @@ async def test_embedding_rate_limiting_realistic():
 
     # Clear the config and rate limiter caches to ensure our settings are applied
     get_llm_config.cache_clear()
+    EmbeddingRateLimiter.reset_instance()
 
     # Create a fresh config instance and verify settings
     config = get_llm_config()
@@ -166,6 +170,7 @@ async def test_with_mock_failures():
 
     # Clear caches
     get_llm_config.cache_clear()
+    EmbeddingRateLimiter.reset_instance()
 
     # Create a mock engine configured to fail every 3rd request
     engine = MockEmbeddingEngine()

cognee/tests/unit/modules/graph/cognee_graph_elements_test.py

@@ -9,7 +9,7 @@ def test_node_initialization():
     """Test that a Node is initialized correctly."""
     node = Node("node1", {"attr1": "value1"}, dimension=2)
     assert node.id == "node1"
-    assert node.attributes == {"attr1": "value1", "vector_distance": 3.5}
+    assert node.attributes == {"attr1": "value1", "vector_distance": np.inf}
     assert len(node.status) == 2
     assert np.all(node.status == 1)
 
@@ -96,7 +96,7 @@ def test_edge_initialization():
     edge = Edge(node1, node2, {"weight": 10}, directed=False, dimension=2)
     assert edge.node1 == node1
     assert edge.node2 == node2
-    assert edge.attributes == {"vector_distance": 3.5, "weight": 10}
+    assert edge.attributes == {"vector_distance": np.inf, "weight": 10}
     assert edge.directed is False
     assert len(edge.status) == 2
     assert np.all(edge.status == 1)