NebulaGraph MCP Server

Verified
Apache 2.0
OverviewInspectSchema Related Servers Reviews Score
tests
import os
import time

import pytest
from dotenv import load_dotenv

# Load environment variables
load_dotenv()


# Ensure correct environment variable settings
@pytest.fixture(scope="session", autouse=True)
def setup_env():
    os.environ["NEBULA_VERSION"] = "v3"
    os.environ["NEBULA_HOST"] = os.getenv("NEBULA_HOST", "127.0.0.1")
    os.environ["NEBULA_PORT"] = os.getenv("NEBULA_PORT", "9669")
    os.environ["NEBULA_USER"] = os.getenv("NEBULA_USER", "root")
    os.environ["NEBULA_PASSWORD"] = os.getenv("NEBULA_PASSWORD", "nebula")


# Setup commands for test data
SETUP_COMMANDS = [
    # Group 1: Create space
    ["CREATE SPACE IF NOT EXISTS test_graph(vid_type=FIXED_STRING(30))"],
    # Group 2: Use space and create tags and edge types
    [
        "USE test_graph",
        "CREATE TAG IF NOT EXISTS person(name string, age int)",
        "CREATE EDGE IF NOT EXISTS knows(years int)",
        "CREATE EDGE IF NOT EXISTS reports_to(department string)",
    ],
    # Group 3: Insert vertices
    [
        "USE test_graph",
        'INSERT VERTEX person(name, age) VALUES "person1":("Alice", 30)',
        'INSERT VERTEX person(name, age) VALUES "person2":("Bob", 32)',
        'INSERT VERTEX person(name, age) VALUES "person3":("Charlie", 45)',
        'INSERT VERTEX person(name, age) VALUES "person4":("David", 50)',
        'INSERT VERTEX person(name, age) VALUES "person5":("Eve", 55)',
    ],
    # Group 4: Insert edges
    [
        "USE test_graph",
        'INSERT EDGE knows(years) VALUES "person1"->"person2":(5)',
        'INSERT EDGE knows(years) VALUES "person2"->"person3":(3)',
        'INSERT EDGE knows(years) VALUES "person3"->"person5":(10)',
        'INSERT EDGE reports_to(department) VALUES "person1"->"person3":("Engineering")',
        'INSERT EDGE reports_to(department) VALUES "person2"->"person3":("Engineering")',
        'INSERT EDGE reports_to(department) VALUES "person3"->"person4":("Management")',
        'INSERT EDGE reports_to(department) VALUES "person4"->"person5":("Executive")',
    ],
]

# Skip integration tests (unless explicitly enabled)
pytestmark = pytest.mark.skipif(
    os.getenv("RUN_INTEGRATION_TESTS") != "true",
    reason="Integration tests require a running NebulaGraph instance and specific test data",
)


def execute_with_retry(session, command, max_retries=3, retry_interval=2):
    """Execute command and retry on failure"""
    for attempt in range(max_retries):
        try:
            print(f"Executing: {command}")
            result = session.execute(command)
            if result.is_succeeded():
                return result
            else:
                print(
                    f"Error executing command (attempt {attempt + 1}/{max_retries}): {result.error_msg()}"
                )
                if attempt < max_retries - 1:
                    time.sleep(retry_interval)
        except Exception as e:
            print(
                f"Error executing command (attempt {attempt + 1}/{max_retries}): {e!s}"
            )
            if attempt < max_retries - 1:
                time.sleep(retry_interval)

    # final attempt
    print(f"Final attempt to execute: {command}")
    return session.execute(command)


@pytest.fixture(scope="module", autouse=True)
def setup_connection():
    """Initialize connection pool and prepare test data"""
    from nebulagraph_mcp_server.server import config, global_pool

    # Initialize connection
    global_pool.init(
        [
            (
                os.getenv("NEBULA_HOST", "127.0.0.1"),
                int(os.getenv("NEBULA_PORT", "9669")),
            )
        ],
        config,
    )

    # If running integration tests, create test data
    if os.getenv("RUN_INTEGRATION_TESTS") == "true":
        session = global_pool.get_session(
            os.getenv("NEBULA_USER", "root"), os.getenv("NEBULA_PASSWORD", "nebula")
        )

        try:
            # Execute command groups
            for group_index, command_group in enumerate(SETUP_COMMANDS):
                print(
                    f"\nExecuting command group {group_index + 1}/{len(SETUP_COMMANDS)}:"
                )

                # Execute commands in the group
                for command in command_group:
                    if command.strip():
                        execute_with_retry(session, command.strip())

                # Add appropriate wait time after group type
                if group_index == 0:  # After creating space
                    print("Waiting for space creation to complete (30 seconds)...")
                    time.sleep(30)
                elif group_index == 1:  # After creating tags and edge types
                    print("Waiting for metadata synchronization (10 seconds)...")
                    time.sleep(10)
                else:  # Short wait after other operations
                    print("Short wait (2 seconds)...")
                    time.sleep(2)

            # Final verification
            print("\nVerifying data creation:")
            verify_commands = [
                "USE test_graph",
                "SHOW TAGS",
                "SHOW EDGES",
                "MATCH (n:person) RETURN count(n)",
            ]

            for command in verify_commands:
                result = execute_with_retry(session, command)
                print(
                    f"Verification result: {command} - {'Success' if result.is_succeeded() else 'Failed'}"
                )

        finally:
            session.release()

    yield

    # Clean up connection
    global_pool.close()


# Integration test: get all spaces
@pytest.mark.asyncio
async def test_integration_list_spaces():
    """Test the list spaces function"""
    from nebulagraph_mcp_server.server import list_spaces

    # Call the list_spaces tool function
    result = list_spaces()

    # Validate the result
    assert "Available spaces" in result
    assert "test_graph" in result  # Assuming the test space is created


# Integration test: execute query
@pytest.mark.asyncio
async def test_integration_execute_query():
    """Test the execute query function"""
    from nebulagraph_mcp_server.server import execute_query

    # Execute the query
    result = execute_query(
        "MATCH (n:person) RETURN id(n), n.name, n.age LIMIT 3",
        "test_graph",
    )

    # Validate the result
    assert "Results" in result
    assert "id(n) | n.name | n.age" in result


# Integration test: find path
@pytest.mark.asyncio
async def test_integration_find_path():
    """Test the find path function"""
    from nebulagraph_mcp_server.server import find_path

    # Test finding paths from person1 to person5
    result = find_path("person1", "person5", "test_graph", 3, 10)
    # Validate the result
    print(f"Find path result: {result}")  # Debug
    assert "Find paths from person1 to person5" in result
    # Since the test graph, person1 to person5 needs at least 3 steps,
    # when depth=3, there should be a path, and the result should contain path information
    # If no path is found, the result should contain "No paths found"
    assert "Path" in result or "No paths found" in result
    # Test non-existent path
    result2 = find_path("person1", "nonexistent", "test_graph", 3, 10)
    assert "No paths found" in result2 or "Query failed" in result2


# Integration test: get neighbors
@pytest.mark.asyncio
async def test_integration_get_neighbors():
    """Test the get neighbors function"""
    from nebulagraph_mcp_server.server import find_neighbors

    # Test getting neighbors of person1
    result = find_neighbors("person1", "test_graph", 1)
    # Validate the result
    assert "Vertex person1 neighbors (depth 1):" in result
    assert "person2" in result
    assert "person3" in result
    assert "person4" not in result
    assert "person5" not in result