r2r/tests/integration/test_graphs.py

import uuid

import pytest

from r2r import R2RClient, R2RException


@pytest.fixture(scope="session")
def config():
    class TestConfig:
        base_url = "http://localhost:7272"
        superuser_email = "admin@example.com"
        superuser_password = "change_me_immediately"

    return TestConfig()


@pytest.fixture(scope="session")
def client(config):
    """Create a client instance and possibly log in as a superuser."""
    client = R2RClient(config.base_url)
    client.users.login(config.superuser_email, config.superuser_password)
    return client


@pytest.fixture
def test_collection(client):
    """Create a test collection (and thus a graph) for testing, then delete it afterwards."""
    resp = client.collections.create(
        name=f"Test Collection {uuid.uuid4()}",
        description="A sample collection for graph tests",
    )
    collection_id = resp["results"]["id"]
    yield collection_id
    # Cleanup if needed
    # If there's a deletion endpoint for collections, call it here.
    client.collections.delete(id=collection_id)


def test_list_graphs(client):
    resp = client.graphs.list(limit=5)
    assert "results" in resp, "No results field in list response"


def test_create_and_get_graph(client, test_collection):
    # `test_collection` fixture creates a collection and returns ID
    collection_id = test_collection
    resp = client.graphs.retrieve(collection_id=collection_id)["results"]
    assert resp["collection_id"] == collection_id, "Graph ID mismatch"


def test_update_graph(client, test_collection):
    collection_id = test_collection
    new_name = "Updated Test Graph Name"
    new_description = "Updated test description"

    resp = client.graphs.update(
        collection_id=collection_id, name=new_name, description=new_description
    )["results"]

    assert resp["name"] == new_name, "Name not updated correctly"
    assert (
        resp["description"] == new_description
    ), "Description not updated correctly"


def test_list_entities(client, test_collection):
    collection_id = test_collection
    resp = client.graphs.list_entities(collection_id=collection_id, limit=5)[
        "results"
    ]
    assert isinstance(resp, list), "No results array in entities response"


def test_create_and_get_entity(client, test_collection):
    collection_id = test_collection
    entity_name = "Test Entity"
    entity_description = "Test entity description"

    create_resp = client.graphs.create_entity(
        collection_id=collection_id,
        name=entity_name,
        description=entity_description,
    )["results"]
    entity_id = create_resp["id"]

    resp = client.graphs.get_entity(
        collection_id=collection_id, entity_id=entity_id
    )["results"]
    assert resp["name"] == entity_name, "Entity name mismatch"


def test_list_relationships(client, test_collection):
    collection_id = test_collection
    resp = client.graphs.list_relationships(
        collection_id=collection_id, limit=5
    )["results"]
    assert isinstance(resp, list), "No results array in relationships response"


def test_create_and_get_relationship(client, test_collection):
    collection_id = test_collection

    # Create two entities
    entity1 = client.graphs.create_entity(
        collection_id=collection_id,
        name="Entity 1",
        description="Entity 1 description",
    )["results"]
    entity2 = client.graphs.create_entity(
        collection_id=collection_id,
        name="Entity 2",
        description="Entity 2 description",
    )["results"]

    # Create relationship
    rel_resp = client.graphs.create_relationship(
        collection_id=collection_id,
        subject="Entity 1",
        subject_id=entity1["id"],
        predicate="related_to",
        object="Entity 2",
        object_id=entity2["id"],
        description="Test relationship",
    )["results"]
    relationship_id = rel_resp["id"]

    # Get relationship
    resp = client.graphs.get_relationship(
        collection_id=collection_id, relationship_id=relationship_id
    )["results"]
    assert resp["predicate"] == "related_to", "Relationship predicate mismatch"


def test_build_communities(client, test_collection):
    collection_id = test_collection

    # Create two entities
    entity1 = client.graphs.create_entity(
        collection_id=collection_id,
        name="Entity 1",
        description="Entity 1 description",
    )["results"]
    entity2 = client.graphs.create_entity(
        collection_id=collection_id,
        name="Entity 2",
        description="Entity 2 description",
    )["results"]

    # Create relationship
    rel_resp = client.graphs.create_relationship(
        collection_id=collection_id,
        subject="Entity 1",
        subject_id=entity1["id"],
        predicate="related_to",
        object="Entity 2",
        object_id=entity2["id"],
        description="Test relationship",
    )["results"]
    relationship_id = rel_resp["id"]

    # Build communities
    # Adjust parameters as needed if `run_type` and `settings` differ.
    # The router expects `run_type` and `graph_enrichment_settings`.
    resp = client.graphs.build(
        collection_id=collection_id,
        run_type="run",
        # graph_enrichment_settings={"use_semantic_clustering": True},
        run_with_orchestration=False,
    )["results"]

    # After building, list communities
    resp = client.graphs.list_communities(
        collection_id=collection_id, limit=5
    )["results"]
    # We cannot guarantee communities are created if no entities or special conditions apply.
    # If no communities, we may skip this assert or ensure at least no error occurred.
    assert isinstance(resp, list), "No communities array returned."


def test_list_communities(client, test_collection):
    collection_id = test_collection
    resp = client.graphs.list_communities(
        collection_id=collection_id, limit=5
    )["results"]
    assert isinstance(resp, list), "No results array in communities response"


def test_create_and_get_community(client, test_collection):
    collection_id = test_collection
    community_name = "Test Community"
    community_summary = "Test community summary"

    create_resp = client.graphs.create_community(
        collection_id=collection_id,
        name=community_name,
        summary=community_summary,
        findings=["Finding 1", "Finding 2"],
        rating=8,
    )["results"]
    community_id = create_resp["id"]

    resp = client.graphs.get_community(
        collection_id=collection_id, community_id=community_id
    )["results"]
    assert resp["name"] == community_name, "Community name mismatch"


def test_update_community(client, test_collection):
    collection_id = test_collection
    # Create a community to update
    create_resp = client.graphs.create_community(
        collection_id=collection_id,
        name="Community to update",
        summary="Original summary",
        findings=["Original finding"],
        rating=7,
    )["results"]
    community_id = create_resp["id"]

    # Update the community
    resp = client.graphs.update_community(
        collection_id=collection_id,
        community_id=community_id,
        name="Updated Community",
        summary="Updated summary",
        findings=["New finding"],
        rating=9,
    )["results"]

    assert resp["name"] == "Updated Community", "Community update failed"


def test_pull_operation(client, test_collection):
    collection_id = test_collection
    resp = client.graphs.pull(collection_id=collection_id)["results"]
    # Depending on your system, the pull might require documents in the collection to actually do something.
    # If no documents are present, it might return success=False or a warning. Check behavior and adjust test.
    assert "success" in resp, "No success indicator in pull response"


def test_error_handling(client):
    # Test retrieving a graph with invalid ID
    invalid_id = "not-a-uuid"
    with pytest.raises(R2RException) as exc_info:
        client.graphs.retrieve(collection_id=invalid_id)
    # Expecting a 422 or 404 error. Adjust as per your API's expected response.
    assert exc_info.value.status_code in [
        400,
        422,
        404,
    ], "Expected an error for invalid ID."