# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for the amazon-neptune MCP Server."""
import pytest
from awslabs.amazon_neptune_mcp_server.server import (
get_graph,
get_schema,
get_schema_resource,
get_status,
get_status_resource,
main,
run_gremlin_query,
run_opencypher_query,
)
from unittest.mock import MagicMock, patch
@pytest.mark.asyncio
class TestServerTools:
"""Test class for server tool functions that interact with the Neptune graph."""
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_get_status(self, mock_get_graph):
"""Test that get_status correctly returns the status from the graph.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The status method is called on the graph instance
3. The result from the graph's status method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_graph.status.return_value = 'Connected'
mock_get_graph.return_value = mock_graph
# Act
result = get_status()
# Assert
assert result == 'Connected'
mock_graph.status.assert_called_once()
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_get_schema(self, mock_get_graph):
"""Test that get_schema correctly returns the schema from the graph.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The schema method is called on the graph instance
3. The result from the graph's schema method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_schema = MagicMock()
mock_graph.schema.return_value = mock_schema
mock_get_graph.return_value = mock_graph
# Act
result = get_schema()
# Assert
assert result == mock_schema
mock_graph.schema.assert_called_once()
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_run_opencypher_query(self, mock_get_graph):
"""Test that run_opencypher_query correctly executes a query without parameters.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The query_opencypher method is called with the correct query and None parameters
3. The result from the graph's query_opencypher method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_result = {'results': [{'n': {'id': '1'}}]}
mock_graph.query_opencypher.return_value = mock_result
mock_get_graph.return_value = mock_graph
# Act
result = run_opencypher_query('MATCH (n) RETURN n LIMIT 1')
# Assert
assert result == mock_result
mock_graph.query_opencypher.assert_called_once_with('MATCH (n) RETURN n LIMIT 1', None)
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_run_opencypher_query_with_parameters(self, mock_get_graph):
"""Test that run_opencypher_query correctly executes a query with parameters.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The query_opencypher method is called with the correct query and parameters
3. The result from the graph's query_opencypher method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_result = {'results': [{'n': {'id': '1'}}]}
mock_graph.query_opencypher.return_value = mock_result
mock_get_graph.return_value = mock_graph
parameters = {'id': '1'}
# Act
result = run_opencypher_query('MATCH (n) WHERE n.id = $id RETURN n', parameters)
# Assert
assert result == mock_result
mock_graph.query_opencypher.assert_called_once_with(
'MATCH (n) WHERE n.id = $id RETURN n', parameters
)
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_run_gremlin_query(self, mock_get_graph):
"""Test that run_gremlin_query correctly executes a Gremlin query.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The query_gremlin method is called with the correct query
3. The result from the graph's query_gremlin method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_result = {'results': [{'id': '1'}]}
mock_graph.query_gremlin.return_value = mock_result
mock_get_graph.return_value = mock_graph
# Act
result = run_gremlin_query('g.V().limit(1)')
# Assert
assert result == mock_result
mock_graph.query_gremlin.assert_called_once_with('g.V().limit(1)')
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_get_status_resource(self, mock_get_graph):
"""Test that get_status_resource correctly returns the status from the graph.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The status method is called on the graph instance
3. The result from the graph's status method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_graph.status.return_value = 'AVAILABLE'
mock_get_graph.return_value = mock_graph
# Act
result = get_status_resource()
# Assert
assert result == 'AVAILABLE'
mock_graph.status.assert_called_once()
@patch('awslabs.amazon_neptune_mcp_server.server.get_graph')
async def test_get_schema_resource(self, mock_get_graph):
"""Test that get_schema_resource correctly returns the schema from the graph.
This test verifies that:
1. The get_graph function is called to obtain the graph instance
2. The schema method is called on the graph instance
3. The result from the graph's schema method is returned unchanged.
"""
# Arrange
mock_graph = MagicMock()
mock_schema = MagicMock()
mock_graph.schema.return_value = mock_schema
mock_get_graph.return_value = mock_graph
# Act
result = get_schema_resource()
# Assert
assert result == mock_schema
mock_graph.schema.assert_called_once()
@pytest.mark.asyncio
class TestGraphInitialization:
"""Test class for the graph initialization functionality."""
@patch('os.environ.get')
@patch('awslabs.amazon_neptune_mcp_server.server.NeptuneServer')
async def test_get_graph_initialization(self, mock_neptune_server, mock_environ_get):
"""Test that get_graph correctly initializes a NeptuneServer instance.
This test verifies that:
1. Environment variables are correctly read
2. NeptuneServer is initialized with the correct parameters
3. The same instance is returned on subsequent calls (singleton pattern).
"""
# Arrange
mock_environ_get.side_effect = lambda key, default=None: {
'NEPTUNE_ENDPOINT': 'neptune-db://test-endpoint',
'NEPTUNE_USE_HTTPS': 'True',
}.get(key, default)
mock_server = MagicMock()
mock_neptune_server.return_value = mock_server
# Act
graph = get_graph()
# Assert
assert graph == mock_server
mock_neptune_server.assert_called_once_with('neptune-db://test-endpoint', use_https=True)
# Call again to verify singleton behavior
graph2 = get_graph()
assert graph2 == graph
mock_neptune_server.assert_called_once() # Should not be called again
@patch('os.environ.get')
async def test_get_graph_missing_endpoint(self, mock_environ_get):
"""Test that get_graph raises an error when the NEPTUNE_ENDPOINT environment variable is missing.
This test verifies that:
1. When NEPTUNE_ENDPOINT is None, a ValueError is raised
2. The error message correctly indicates the missing environment variable.
"""
# Arrange
mock_environ_get.side_effect = lambda key, default=None: {
'NEPTUNE_ENDPOINT': None,
'NEPTUNE_USE_HTTPS': 'True',
}.get(key, default)
# Reset the global _graph variable
import awslabs.amazon_neptune_mcp_server.server
awslabs.amazon_neptune_mcp_server.server._graph = None
# Act & Assert
with pytest.raises(ValueError, match='NEPTUNE_ENDPOINT environment variable is not set'):
get_graph()
@patch('os.environ.get')
@patch('awslabs.amazon_neptune_mcp_server.server.NeptuneServer')
async def test_get_graph_with_https_false(self, mock_neptune_server, mock_environ_get):
"""Test that get_graph correctly handles HTTPS settings from environment variables.
This test verifies that:
1. When NEPTUNE_USE_HTTPS is set to "false", use_https is set to False
2. NeptuneServer is initialized with the correct parameters.
"""
# Arrange
mock_environ_get.side_effect = lambda key, default=None: {
'NEPTUNE_ENDPOINT': 'neptune-db://test-endpoint',
'NEPTUNE_USE_HTTPS': 'false',
}.get(key, default)
# Reset the global _graph variable
import awslabs.amazon_neptune_mcp_server.server
awslabs.amazon_neptune_mcp_server.server._graph = None
mock_server = MagicMock()
mock_neptune_server.return_value = mock_server
# Act
graph = get_graph()
# Assert
assert graph == mock_server
mock_neptune_server.assert_called_once_with('neptune-db://test-endpoint', use_https=False)
@pytest.mark.asyncio
class TestMainFunction:
"""Test class for the main function that runs the MCP server."""
@patch('awslabs.amazon_neptune_mcp_server.server.mcp')
async def test_main_default(self, mock_mcp):
"""Test that main correctly runs the server with default settings."""
# Arrange
# Act
main()
# Assert
assert mock_mcp.run.call_count == 1
