From 41b69083a316aa8a4874a4acf5cea87f9a7cd489 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Florentin=20D=C3=B6rre?= <florentin.dorre@neo4j.com>
Date: Wed, 24 Apr 2024 16:55:22 +0200
Subject: [PATCH] Add instances/sizing endpoint

---
 aura/instances/__init__.py          |   3 +
 aura/instances/sizing.py            |  53 +++++++++++++
 tests/unit/instances/test_sizing.py | 117 ++++++++++++++++++++++++++++
 3 files changed, 173 insertions(+)
 create mode 100644 aura/instances/sizing.py
 create mode 100644 tests/unit/instances/test_sizing.py

diff --git a/aura/instances/__init__.py b/aura/instances/__init__.py
index cbc635d..744e5fd 100644
--- a/aura/instances/__init__.py
+++ b/aura/instances/__init__.py
@@ -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
@@ -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)
diff --git a/aura/instances/sizing.py b/aura/instances/sizing.py
new file mode 100644
index 0000000..186bc57
--- /dev/null
+++ b/aura/instances/sizing.py
@@ -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))
diff --git a/tests/unit/instances/test_sizing.py b/tests/unit/instances/test_sizing.py
new file mode 100644
index 0000000..a0a4f93
--- /dev/null
+++ b/tests/unit/instances/test_sizing.py
@@ -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,
+    )
\ No newline at end of file