cognee 0.5.1.dev0__py3-none-any.whl → 0.5.2__py3-none-any.whl

cognee/tests/unit/api/test_get_raw_data_endpoint.py

@@ -0,0 +1,206 @@
+import io
+import uuid
+from contextlib import asynccontextmanager
+from types import SimpleNamespace
+from unittest.mock import AsyncMock
+
+import pytest
+from fastapi import FastAPI
+from fastapi.testclient import TestClient
+
+from cognee.modules.users.methods import get_authenticated_user
+
+
+@pytest.fixture(scope="session")
+def test_client():
+    from cognee.api.v1.datasets.routers.get_datasets_router import get_datasets_router
+
+    app = FastAPI()
+    app.include_router(get_datasets_router(), prefix="/api/v1/datasets")
+
+    with TestClient(app) as c:
+        yield c
+
+
+@pytest.fixture
+def client(test_client):
+    async def override_get_authenticated_user():
+        return SimpleNamespace(
+            id=str(uuid.uuid4()),
+            email="default@example.com",
+            is_active=True,
+            tenant_id=str(uuid.uuid4()),
+        )
+
+    import importlib
+
+    datasets_router_module = importlib.import_module(
+        "cognee.api.v1.datasets.routers.get_datasets_router"
+    )
+    datasets_router_module.send_telemetry = lambda *args, **kwargs: None
+
+    test_client.app.dependency_overrides[get_authenticated_user] = override_get_authenticated_user
+    yield test_client
+    test_client.app.dependency_overrides.pop(get_authenticated_user, None)
+
+
+def _patch_raw_download_dependencies(
+    monkeypatch, *, dataset_id, data_id, raw_data_location, name, mime_type
+):
+    """
+    Patch the internal dataset/data lookups used by GET /datasets/{dataset_id}/data/{data_id}/raw.
+    Keeps the test focused on response behavior (FileResponse vs StreamingResponse).
+    """
+    import importlib
+
+    datasets_router_module = importlib.import_module(
+        "cognee.api.v1.datasets.routers.get_datasets_router"
+    )
+
+    monkeypatch.setattr(
+        datasets_router_module,
+        "get_authorized_existing_datasets",
+        AsyncMock(return_value=[SimpleNamespace(id=dataset_id)]),
+    )
+
+    import cognee.modules.data.methods as data_methods_module
+
+    monkeypatch.setattr(
+        data_methods_module,
+        "get_dataset_data",
+        AsyncMock(return_value=[SimpleNamespace(id=data_id)]),
+    )
+    monkeypatch.setattr(
+        data_methods_module,
+        "get_data",
+        AsyncMock(
+            return_value=SimpleNamespace(
+                id=data_id,
+                raw_data_location=raw_data_location,
+                name=name,
+                mime_type=mime_type,
+            )
+        ),
+    )
+
+
+def test_get_raw_data_local_file_downloads_bytes(client, monkeypatch, tmp_path):
+    """Downloads bytes from a file:// raw_data_location."""
+    dataset_id = uuid.uuid4()
+    data_id = uuid.uuid4()
+
+    file_path = tmp_path / "example.txt"
+    content = b"hello from disk"
+    file_path.write_bytes(content)
+
+    _patch_raw_download_dependencies(
+        monkeypatch,
+        dataset_id=dataset_id,
+        data_id=data_id,
+        raw_data_location=file_path.as_uri(),
+        name="example.txt",
+        mime_type="text/plain",
+    )
+
+    response = client.get(f"/api/v1/datasets/{dataset_id}/data/{data_id}/raw")
+    assert response.status_code == 200
+    assert response.content == content
+
+
+def test_get_raw_data_s3_streams_bytes_without_s3_dependency(client, monkeypatch):
+    """Streams bytes from an s3:// raw_data_location (mocked)."""
+    dataset_id = uuid.uuid4()
+    data_id = uuid.uuid4()
+
+    _patch_raw_download_dependencies(
+        monkeypatch,
+        dataset_id=dataset_id,
+        data_id=data_id,
+        raw_data_location="s3://bucket/path/to/file.txt",
+        name="file.txt",
+        mime_type="text/plain",
+    )
+
+    import cognee.infrastructure.files.utils.open_data_file as open_data_file_module
+
+    @asynccontextmanager
+    async def fake_open_data_file(_file_path: str, mode: str = "rb", **_kwargs):
+        assert mode == "rb"
+        yield io.BytesIO(b"hello from s3")
+
+    monkeypatch.setattr(open_data_file_module, "open_data_file", fake_open_data_file)
+
+    response = client.get(f"/api/v1/datasets/{dataset_id}/data/{data_id}/raw")
+    assert response.status_code == 200
+    assert response.content == b"hello from s3"
+    assert response.headers.get("content-disposition") == 'attachment; filename="file.txt"'
+
+
+def test_get_raw_data_unsupported_scheme_returns_501(client, monkeypatch):
+    """Returns 501 for unsupported raw_data_location schemes (e.g., http://)."""
+    dataset_id = uuid.uuid4()
+    data_id = uuid.uuid4()
+
+    _patch_raw_download_dependencies(
+        monkeypatch,
+        dataset_id=dataset_id,
+        data_id=data_id,
+        raw_data_location="http://example.com/file.txt",
+        name="file.txt",
+        mime_type="text/plain",
+    )
+
+    response = client.get(f"/api/v1/datasets/{dataset_id}/data/{data_id}/raw")
+    assert response.status_code == 501
+    assert "Storage scheme 'http' not supported" in response.json()["detail"]
+
+
+def test_get_raw_data_plain_path_downloads_bytes(client, monkeypatch, tmp_path):
+    """Downloads bytes from a plain local path (no scheme)."""
+    dataset_id = uuid.uuid4()
+    data_id = uuid.uuid4()
+
+    file_path = tmp_path / "plain.txt"
+    content = b"plain content"
+    file_path.write_bytes(content)
+
+    _patch_raw_download_dependencies(
+        monkeypatch,
+        dataset_id=dataset_id,
+        data_id=data_id,
+        raw_data_location=str(file_path),
+        name="plain.txt",
+        mime_type="text/plain",
+    )
+
+    response = client.get(f"/api/v1/datasets/{dataset_id}/data/{data_id}/raw")
+    assert response.status_code == 200
+    assert response.content == content
+
+
+def test_get_raw_data_encoded_path_downloads_bytes(client, monkeypatch, tmp_path):
+    """Downloads bytes from a percent-encoded file URI (e.g. spaces)."""
+    dataset_id = uuid.uuid4()
+    data_id = uuid.uuid4()
+
+    file_name = "example file.txt"
+    file_path = tmp_path / file_name
+    content = b"content with spaces"
+    file_path.write_bytes(content)
+
+    # Convert to URI, which should encode space as %20
+    uri = file_path.as_uri()
+    assert "%20" in uri
+
+    _patch_raw_download_dependencies(
+        monkeypatch,
+        dataset_id=dataset_id,
+        data_id=data_id,
+        raw_data_location=uri,
+        name=file_name,
+        mime_type="text/plain",
+    )
+
+    response = client.get(f"/api/v1/datasets/{dataset_id}/data/{data_id}/raw")
+    assert response.status_code == 200
+    assert response.content == content

cognee/tests/unit/eval_framework/answer_generation_test.py

@@ -12,8 +12,9 @@ async def test_answer_generation():
     corpus_list, qa_pairs = DummyAdapter().load_corpus(limit=limit)
 
     mock_retriever = AsyncMock()
-    mock_retriever.get_context = AsyncMock(return_value="Mocked retrieval context")
-    mock_retriever.get_completion = AsyncMock(return_value=["Mocked answer"])
+    mock_retriever.get_retrieved_objects = AsyncMock(return_value=[])
+    mock_retriever.get_context_from_objects = AsyncMock(return_value="Mocked retrieval context")
+    mock_retriever.get_completion_from_context = AsyncMock(return_value=["Mocked answer"])
 
     answer_generator = AnswerGeneratorExecutor()
     answers = await answer_generator.question_answering_non_parallel(
@@ -21,7 +22,7 @@ async def test_answer_generation():
         retriever=mock_retriever,
     )
 
-    mock_retriever.get_context.assert_any_await(qa_pairs[0]["question"])
+    mock_retriever.get_context_from_objects.assert_any_await(qa_pairs[0]["question"], [])
 
     assert len(answers) == len(qa_pairs)
     assert answers[0]["question"] == qa_pairs[0]["question"], (

cognee/tests/unit/infrastructure/databases/cache/test_cache_config.py

@@ -62,6 +62,8 @@ def test_cache_config_to_dict():
         "cache_password": None,
         "agentic_lock_expire": 100,
         "agentic_lock_timeout": 200,
+        "usage_logging": False,
+        "usage_logging_ttl": 604800,
     }
 
 

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": None}
     assert len(node.status) == 2
     assert np.all(node.status == 1)
 
@@ -86,6 +86,46 @@ def test_node_hash():
     assert hash(node) == hash("node1")
 
 
+def test_node_vector_distance_stays_none():
+    """Test that vector_distance remains None when no distances are passed."""
+    node = Node("node1")
+    assert node.attributes.get("vector_distance") is None
+
+    # Verify it stays None even after other operations
+    node.add_attribute("other_attr", "value")
+    assert node.attributes.get("vector_distance") is None
+
+
+def test_node_vector_distance_with_custom_attributes():
+    """Test that vector_distance is None even when node has custom attributes."""
+    node = Node("node1", {"custom": "value", "another": 42})
+    assert node.attributes.get("vector_distance") is None
+    assert node.attributes["custom"] == "value"
+    assert node.attributes["another"] == 42
+
+
+def test_edge_vector_distance_stays_none():
+    """Test that vector_distance remains None when no distances are passed."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2)
+    assert edge.attributes.get("vector_distance") is None
+
+    # Verify it stays None even after other operations
+    edge.add_attribute("other_attr", "value")
+    assert edge.attributes.get("vector_distance") is None
+
+
+def test_edge_vector_distance_with_custom_attributes():
+    """Test that vector_distance is None even when edge has custom attributes."""
+    node1 = Node("node1")
+    node2 = Node("node2")
+    edge = Edge(node1, node2, {"weight": 5, "type": "test"})
+    assert edge.attributes.get("vector_distance") is None
+    assert edge.attributes["weight"] == 5
+    assert edge.attributes["type"] == "test"
+
+
 ### Tests for Edge ###
 
 
@@ -96,7 +136,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": None, "weight": 10}
     assert edge.directed is False
     assert len(edge.status) == 2
     assert np.all(edge.status == 1)