Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add instances/sizing endpoint #26

Open
wants to merge 1 commit into
base: 1.0
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions aura/instances/__init__.py
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
import click

from aura.instances.sizing import size_instance
from .create import create_instance
from .get import get_instance
from .list import list_instances
Expand All @@ -15,6 +17,7 @@ def instances():


instances.add_command(create_instance)
instances.add_command(size_instance)
instances.add_command(get_instance)
instances.add_command(list_instances)
instances.add_command(update_instance)
Expand Down
53 changes: 53 additions & 0 deletions aura/instances/sizing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
import json
from typing import List
import click
from aura.api_command import api_command
from aura.api_repository import make_api_call
from aura.config_repository import CLIConfig
from aura.decorators import pass_config


# pylint: disable=redefined-builtin
@api_command(name="sizing", help_text="Estimate the size of an instance")
@click.option("--type", "-t", prompt=True, help="The instance type")
@click.option(
"--node-count", "-n", prompt=True, help="The estimated node count", type=int
)
@click.option(
"--relationship-count",
"-r",
prompt=True,
help="The estimated relationship count",
type=int,
)
@click.option(
"--algorithm-categories",
"-ac",
help="The GDS algorithm categories to be used",
multiple=True,
default=[],
type=str,
)
@pass_config
def size_instance(
config: CLIConfig,
type: str,
node_count: int,
relationship_count: int,
algorithm_categories: List[str],
):
"""
Estimates the size of an instance using the given options.

Makes "POST /instances/sizing" API request.
"""
path = "/instances/sizing"

data = {
"node_count": node_count,
"relationship_count": relationship_count,
"algorithm_categories": algorithm_categories,
"instance_type": type,
}

return make_api_call("POST", path, data=json.dumps(data))
117 changes: 117 additions & 0 deletions tests/unit/instances/test_sizing.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,117 @@
import pytest
from click.testing import CliRunner
from unittest.mock import Mock
import json

from aura.instances import size_instance
from unit.conftest import printed_data


def mock_response():
mock_res = Mock()
mock_res.status_code = 200
mock_res.json.return_value = {
"data": {
"did_exceed_maximum": False,
"min_required_memory": "1GB",
"recommended_size": "8GB",
}
}

return mock_res


def test_size_instance(api_request, mock_config):
runner = CliRunner()

api_request.return_value = mock_response()

result = runner.invoke(
size_instance,
[
"--type",
"enterprise-ds",
"--node-count",
"42",
"--relationship-count",
"1337",
"--algorithm-categories",
"centrality",
"--algorithm-categories",
"node-embeddings",
],
obj=mock_config,
)

assert result.exit_code == 0
assert result.output == printed_data(
{
"did_exceed_maximum": False,
"min_required_memory": "1GB",
"recommended_size": "8GB"
}
)

api_request.assert_called_once_with(
"POST",
"https://api.neo4j.io/v1/instances/sizing",
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer dummy-token",
},
data=json.dumps(
{
"node_count": 42,
"relationship_count": 1337,
"algorithm_categories": ["centrality", "node-embeddings"],
"instance_type": "enterprise-ds",
}
),
timeout=10,
)


def test_sizing_instance_without_ac(api_request, mock_config):
runner = CliRunner()

api_request.return_value = mock_response()

result = runner.invoke(
size_instance,
[
"--type",
"enterprise-ds",
"--node-count",
"42",
"--relationship-count",
"1337",
],
obj=mock_config,
)

assert result.exit_code == 0
assert result.output == printed_data(
{
"did_exceed_maximum": False,
"min_required_memory": "1GB",
"recommended_size": "8GB"
}
)

api_request.assert_called_once_with(
"POST",
"https://api.neo4j.io/v1/instances/sizing",
headers={
"Content-Type": "application/json",
"Authorization": f"Bearer dummy-token",
},
data=json.dumps(
{
"node_count": 42,
"relationship_count": 1337,
"algorithm_categories": [],
"instance_type": "enterprise-ds",
}
),
timeout=10,
)