Skip to content

Commit 2ce793d

Browse files
committed
add cloud provider and region validation, tests
tested with unit tests and claude desktop
1 parent 3b53ba8 commit 2ce793d

File tree

5 files changed

+91
-8
lines changed

5 files changed

+91
-8
lines changed

servers/mcp-neo4j-cloud-aura-api/src/mcp_neo4j_aura_manager/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
33
import argparse
44
import os
55
import logging
6+
import sys
7+
68

79
# Configure logging
810
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')

servers/mcp-neo4j-cloud-aura-api/src/mcp_neo4j_aura_manager/server.py

Lines changed: 34 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,5 @@
1-
import argparse
2-
import asyncio
31
import json
42
import logging
5-
import os
6-
import sys
73
import time
84
from typing import Any, Dict, List, Optional, Union
95

@@ -19,6 +15,30 @@
1915
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
2016
logger = logging.getLogger(__name__)
2117

18+
19+
def _validate_region(cloud_provider: str, region: str) -> None:
20+
"""
21+
Validate the region exists for the given cloud provider.
22+
23+
Args:
24+
cloud_provider: The cloud provider to validate the region for
25+
region: The region to validate
26+
27+
Returns:
28+
None
29+
30+
Raises:
31+
ValueError: If the region is not valid for the given cloud provider
32+
"""
33+
34+
if cloud_provider == "gcp" and region.count("-") != 1:
35+
raise ValueError(f"Invalid region for GCP: {region}. Must follow the format 'region-zonenumber'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
36+
elif cloud_provider == "aws" and region.count("-") != 2:
37+
raise ValueError(f"Invalid region for AWS: {region}. Must follow the format 'region-zone-number'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
38+
elif cloud_provider == "azure" and region.count("-") != 0:
39+
raise ValueError(f"Invalid region for Azure: {region}. Must follow the format 'regionzone'. Refer to https://neo4j.com/docs/aura/managing-instances/regions/ for valid regions.")
40+
41+
2242
class AuraAPIClient:
2343
"""Client for interacting with Neo4j Aura API."""
2444

@@ -172,12 +192,17 @@ def create_instance(self, tenant_id: str, name: str, memory: int = 1, region: st
172192
if cloud_provider and cloud_provider not in ["gcp", "aws", "azure"]:
173193
raise ValueError("cloud_provider must be one of: gcp, aws, azure")
174194

195+
if vector_optimized and memory < 4:
196+
raise ValueError("vector optimized instances must have at least 4GB memory")
197+
175198
# If cloning, source_instance_id is required
176199
if source_instance_id is not None:
177200
if not isinstance(source_instance_id, str):
178201
raise ValueError("source_instance for clone from instance must be defined")
179202
else:
180203
raise ValueError(f"Invalid type {type}")
204+
205+
_validate_region(cloud_provider, region)
181206

182207
url = f"{self.BASE_URL}/instances"
183208
payload = {
@@ -222,6 +247,9 @@ def update_instance(self, instance_id: str, name: Optional[str] = None,
222247
if vector_optimized is not None:
223248
payload["vector_optimized"] = str(vector_optimized).lower()
224249

250+
if payload["vector_optimized"] == "true" and int(payload["memory"]) < 4:
251+
raise ValueError("vector optimized instances must have at least 4GB memory")
252+
225253
print("Update instance payload:")
226254
print(payload)
227255
response = requests.patch(url, headers=self._get_headers(), json=payload)
@@ -485,7 +513,7 @@ async def handle_list_tools() -> List[types.Tool]:
485513
},
486514
"region": {
487515
"type": "string",
488-
"description": "Region for the instance (e.g., 'us-east-1')",
516+
"description": "Region for the instance (e.g., 'us-central1')",
489517
"default": "us-central1"
490518
},
491519
"type": {
@@ -495,7 +523,7 @@ async def handle_list_tools() -> List[types.Tool]:
495523
},
496524
"vector_optimized": {
497525
"type": "boolean",
498-
"description": "Whether the instance is optimized for vector operations",
526+
"description": "Whether the instance is optimized for vector operations. Only allowed for instance with more than 4GB memory.",
499527
"default": False
500528
},
501529
"cloud_provider": {
Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
uv run pytest tests/test_aura_manager.py
21
if [ -f .env ]; then
3-
uv run --env-file .env pytest tests/test_aura_integration.py
2+
uv run --env-file .env pytest tests
3+
else
4+
uv run pytest tests/test_aura_manager.py
45
fi

servers/mcp-neo4j-cloud-aura-api/tests/test_aura_integration.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -263,3 +263,11 @@ def test_create_and_delete_instance_integration(aura_client, test_type):
263263
assert "status" in instance_details
264264
print(f"Deleted test instance {instance_id}: {delete_result} {instance_details}")
265265

266+
267+
def test_create_instance_vector_optimized_and_memory_less_than_4_should_raise_error(aura_client):
268+
with pytest.raises(ValueError):
269+
aura_client.create_instance(memory=3, vector_optimized=True, tenant_id="test-tenant-1", name="Test Instance")
270+
271+
def test_update_instance_vector_optimized_and_memory_less_than_4_should_raise_error(aura_client):
272+
with pytest.raises(ValueError):
273+
aura_client.update_instance(instance_id="test-instance-1", memory=3, vector_optimized=True)
Lines changed: 44 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,44 @@
1+
from mcp_neo4j_aura_manager.server import _validate_region
2+
import pytest
3+
4+
def test_validate_region_aws_valid():
5+
# Test GCP regions
6+
assert _validate_region("aws", "us-east-1") is None
7+
assert _validate_region("aws", "eu-west-1") is None
8+
assert _validate_region("aws", "eu-east-1") is None
9+
10+
def test_validate_region_aws_invalid():
11+
# Test GCP regions
12+
with pytest.raises(ValueError):
13+
_validate_region("aws", "us-east1")
14+
with pytest.raises(ValueError):
15+
_validate_region("aws", "euwest")
16+
with pytest.raises(ValueError):
17+
_validate_region("aws", "eu-west-1-1-1")
18+
19+
def test_validate_region_gcp_valid():
20+
# Test GCP regions
21+
assert _validate_region("gcp", "us-central1") is None
22+
assert _validate_region("gcp", "europe-west1") is None
23+
assert _validate_region("gcp", "us-central2") is None
24+
25+
def test_validate_region_gcp_invalid():
26+
# Test GCP regions
27+
with pytest.raises(ValueError):
28+
_validate_region("gcp", "us-east-1")
29+
with pytest.raises(ValueError):
30+
_validate_region("gcp", "eu-west-1-1")
31+
with pytest.raises(ValueError):
32+
_validate_region("gcp", "euwest")
33+
34+
def test_validate_region_azure_valid():
35+
# Test Azure regions
36+
assert _validate_region("azure", "eastus") is None
37+
assert _validate_region("azure", "northeurope") is None
38+
39+
def test_validate_region_azure_invalid():
40+
# Test Azure regions
41+
with pytest.raises(ValueError):
42+
_validate_region("azure", "us-east-1")
43+
with pytest.raises(ValueError):
44+
_validate_region("azure", "eu-west1")

0 commit comments

Comments
 (0)