Repository: AgnostiqHQ/covalent (language: python)

Issue:
Write unit/functional tests for dask cluster CLI



Discussion / hints:
Please [add your planning poker estimate](https://app.zenhub.com/workspaces/Covalent-Open-Source-627981af7f141b0020c699d1/issues/AgnostiqHQ/covalent/597?planning-poker) with ZenHub @Emmanuel289

Fix patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7edbe3e2c..4d01a4b1d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [UNRELEASED]
 
+### Added
+
+- Dask Cluster CLI functional/unit tests
+
 ### Docs
 
 - Updated RTD concepts, how-to-guides, and api docs with electron dependencies.
diff --git a/covalent_dispatcher/_cli/service.py b/covalent_dispatcher/_cli/service.py
index 1edd1b682..d1307f57f 100644
--- a/covalent_dispatcher/_cli/service.py
+++ b/covalent_dispatcher/_cli/service.py
@@ -27,13 +27,12 @@
 import sys
 import time
 from subprocess import DEVNULL, Popen
-from typing import Optional, Tuple
+from typing import Optional
 
 import click
 import dask.system
 import psutil
-import requests
-from distributed.comm import parse_address, unparse_address
+from distributed.comm import unparse_address
 from distributed.core import connect, rpc
 
 from covalent._shared_files.config import _config_manager as cm


Test patch:
diff --git a/tests/covalent_dispatcher_tests/_cli/service_test.py b/tests/covalent_dispatcher_tests/_cli/service_test.py
index 7b99de4ab..3cafd8a64 100644
--- a/tests/covalent_dispatcher_tests/_cli/service_test.py
+++ b/tests/covalent_dispatcher_tests/_cli/service_test.py
@@ -290,3 +290,190 @@ def test_purge(mocker):
     )
     purge_config_mock.assert_called_once()
     assert result.output == "Covalent server files have been purged.\n"
+
+
+@pytest.mark.parametrize("workers", [1, 2, 3, 4])
+def test_cluster_size(mocker, workers):
+    """
+    Test to ensure that cluster start with the expected number of workers
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = workers
+
+    click = mock.Mock()
+    click.echo = lambda: workers
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_size_mock = mocker.patch(
+        "covalent_dispatcher._cli.service._get_cluster_size", return_value=workers
+    )
+
+    runner = CliRunner()
+    response = runner.invoke(cluster, "--size")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_size_mock.assert_called_once()
+    assert int(response.output) == workers
+
+
+def test_cluster_info(mocker):
+    """
+    Test cluster info CLI
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_info_cli_mock = mocker.patch("covalent_dispatcher._cli.service._get_cluster_info")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--info")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_info_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
+
+
+def test_cluster_status_cli(mocker):
+    """
+    Test cluster status cli
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_status_cli_mock = mocker.patch("covalent_dispatcher._cli.service._get_cluster_status")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--status")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_status_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
+
+
+def test_cluster_address_cli(mocker):
+    """
+    Test cluster address cli
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_cli_mock = mocker.patch("covalent_dispatcher._cli.service._get_cluster_address")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--address")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
+
+
+def test_cluster_logs_cli(mocker):
+    """
+    Test cluster logs cli
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_cli_mock = mocker.patch("covalent_dispatcher._cli.service._get_cluster_logs")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--logs")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
+
+
+def test_cluster_restart_cli(mocker):
+    """
+    Test cluster restart cli
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_cli_mock = mocker.patch("covalent_dispatcher._cli.service._cluster_restart")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--restart")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
+
+
+def test_cluster_scale_cli(mocker):
+    """
+    Test cluster scale cli
+    """
+
+    loop = mock.Mock()
+    loop.run_until_complete.return_value = lambda: ""
+
+    get_event_loop_mock = mocker.patch(
+        "covalent_dispatcher._cli.service.asyncio.get_event_loop", return_value=loop
+    )
+    get_config_mock = mocker.patch("covalent_dispatcher._cli.service.get_config")
+    unparse_addr_mock = mocker.patch("covalent_dispatcher._cli.service.unparse_address")
+    cluster_cli_mock = mocker.patch("covalent_dispatcher._cli.service._cluster_scale")
+    click_echo_mock = mocker.patch("covalent_dispatcher._cli.service.click.echo")
+
+    runner = CliRunner()
+    _ = runner.invoke(cluster, "--scale 1")
+
+    assert get_config_mock.call_count == 2
+    get_event_loop_mock.assert_called_once()
+    unparse_addr_mock.assert_called_once()
+    cluster_cli_mock.assert_called_once()
+    click_echo_mock.assert_called_once()
diff --git a/tests/functional_tests/dask_cluster_cli_test.py b/tests/functional_tests/dask_cluster_cli_test.py
new file mode 100644
index 000000000..445eecb90
--- /dev/null
+++ b/tests/functional_tests/dask_cluster_cli_test.py
@@ -0,0 +1,162 @@
+# Copyright 2021 Agnostiq Inc.
+#
+# This file is part of Covalent.
+#
+# Licensed under the GNU Affero General Public License 3.0 (the "License").
+# A copy of the License may be obtained with this software package or at
+#
+#      https://www.gnu.org/licenses/agpl-3.0.en.html
+#
+# Use of this file is prohibited except in compliance with the License. Any
+# modifications or derivative works of this file must retain this copyright
+# notice, and modified files must contain a notice indicating that they have
+# been altered from the originals.
+#
+# Covalent is distributed in the hope that it will be useful, but WITHOUT
+# ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+# FITNESS FOR A PARTICULAR PURPOSE. See the License for more details.
+#
+# Relief from the License may be granted by purchasing a commercial license.
+
+import ast
+import time
+
+import dask.system
+import pytest
+from click.testing import CliRunner
+from distributed.comm import parse_address
+
+from covalent._shared_files.config import get_config
+from covalent_dispatcher._cli.service import _is_server_running, cluster
+
+
+@pytest.fixture(autouse=True)
+def covalent():
+    """
+    Check to see if covalent is already running in the default
+    configuration
+    """
+    assert _is_server_running()
+
+
+def test_worker_config_option():
+    """
+    Assert that `num_workers` exists as a key in covalent.conf/dask
+    """
+    num_workers = get_config("dask.num_workers")
+    assert int(num_workers) == dask.system.CPU_COUNT
+
+
+def test_cluster_size_cli():
+    """
+    Assert covalent cluster size is equal to the default value
+    """
+    runner = CliRunner()
+    response = runner.invoke(cluster, "--size")
+    assert int(response.output) == dask.system.CPU_COUNT
+
+
+def test_cluster_status_cli():
+    """
+    Assert cluster status CLI for default number of workers in the cluster
+    """
+    runner = CliRunner()
+    response = runner.invoke(cluster, "--status")
+    expected = {}
+    expected["scheduler"] = "running"
+    for i in range(dask.system.CPU_COUNT):
+        expected[f"worker-{i}"] = "running"
+
+    assert expected == ast.literal_eval(response.output)
+
+
+def test_cluster_info_cli():
+    """
+    Test cluster info CLI for default number of workers, memory and threads per
+    worker
+    """
+    runner = CliRunner()
+    response = ast.literal_eval(runner.invoke(cluster, "--info").output)
+    expected_keys = ["type", "id", "address", "services", "started", "workers"]
+    worker_info_expected_keys = [
+        "type",
+        "id",
+        "host",
+        "resources",
+        "local_directory",
+        "name",
+        "nthreads",
+        "memory_limit",
+        "last_seen",
+        "services",
+        "metrics",
+        "status",
+        "nanny",
+    ]
+    assert list(response.keys()) == expected_keys
+    assert len(response["workers"]) == dask.system.CPU_COUNT
+
+    # Check the keys match for each worker in the cluster
+    for _, value in response["workers"].items():
+        assert list(value.keys()) == worker_info_expected_keys
+
+
+def test_cluster_address_cli():
+    """
+    Test cluster address CLI
+    """
+    runner = CliRunner()
+    expected_protocol = "tcp"
+    expected_host = "127.0.0.1"
+
+    response = ast.literal_eval(runner.invoke(cluster, "--address").output)
+    assert list(response.keys()) == ["scheduler", "workers"]
+    assert len(response["workers"]) == dask.system.CPU_COUNT
+
+    scheduler_addr = parse_address(response["scheduler"])
+    scheduler_addr_protocol = scheduler_addr[0]
+    scheduler_addr_host, _ = scheduler_addr[1].split(":")
+    assert scheduler_addr_protocol == expected_protocol
+    assert scheduler_addr_host == expected_host
+
+    for _, worker_addr in response["workers"].items():
+        protocol, host_and_port = parse_address(worker_addr)
+        host, _ = host_and_port.split(":")
+        assert protocol == expected_protocol
+        assert host == expected_host
+
+
+def test_cluster_restart():
+    """
+    Test restarting the cluster by asserting the addresses are different
+    """
+    runner = CliRunner()
+    current_addresses = ast.literal_eval(runner.invoke(cluster, "--address").output)
+
+    response = runner.invoke(cluster, "--restart").output
+    assert response == "Cluster restarted\n"
+
+    new_addresses = ast.literal_eval(runner.invoke(cluster, "--address").output)
+
+    assert current_addresses != new_addresses
+
+
+def test_cluster_scale_up_down():
+    """
+    Test scaling up/down by one worker
+    """
+    target_cluster_size = dask.system.CPU_COUNT + 1
+    runner = CliRunner()
+    response = runner.invoke(cluster, f"--scale {target_cluster_size}")
+    assert response.output == f"Cluster scaled to have {target_cluster_size} workers\n"
+    # Having to wait here for the time it takes to create a new dask worker
+    time.sleep(2)
+    response = runner.invoke(cluster, "--size")
+    assert int(response.output) == target_cluster_size
+
+    # Scale cluster back down
+    response = runner.invoke(cluster, f"--scale {dask.system.CPU_COUNT}")
+    assert response.output == f"Cluster scaled to have {dask.system.CPU_COUNT} workers\n"
+    time.sleep(2)
+    response = runner.invoke(cluster, "--size")
+    assert int(response.output) == dask.system.CPU_COUNT