diff --git a/sunbeam-python/sunbeam/core/common.py b/sunbeam-python/sunbeam/core/common.py index 5b869993..277a5e37 100644 --- a/sunbeam-python/sunbeam/core/common.py +++ b/sunbeam-python/sunbeam/core/common.py @@ -529,3 +529,9 @@ def validate_ip_range(ip_range: str): ipaddress.ip_address(ips[1]) else: raise ValueError("Invalid IP range, must be in the form of 'ip-ip'") + + +def validate_nodes( + ctx: click.core.Context, param: click.core.Option, value: str +) -> list[str]: + return list(set(value.split(","))) diff --git a/sunbeam-python/sunbeam/provider/maintenance/maintenance.py b/sunbeam-python/sunbeam/provider/maintenance/maintenance.py new file mode 100644 index 00000000..b17a8f01 --- /dev/null +++ b/sunbeam-python/sunbeam/provider/maintenance/maintenance.py @@ -0,0 +1,193 @@ +# Copyright (c) 2024 Canonical Ltd. +# +# 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. + +import logging +from typing import Tuple + +import click +from rich.console import Console + +from sunbeam.core.checks import Check, run_preflight_checks +from sunbeam.core.common import ( + BaseStep, + ResultType, + get_step_message, + get_step_result, + run_plan, + validate_nodes, +) +from sunbeam.core.deployment import Deployment +from sunbeam.core.juju import JujuHelper +from sunbeam.provider.local.steps import LocalClusterStatusStep +from sunbeam.provider.maas.steps import MaasClusterStatusStep +from sunbeam.provider.maintenance import checks +from sunbeam.steps.cluster_status import ClusterStatusStep +from sunbeam.steps.hypervisor import EnableHypervisorStep +from sunbeam.steps.maintenance import ( + RunWatcherHostMaintenanceStep, + RunWatcherWorkloadBalancingStep, +) +from sunbeam.utils import click_option_show_hints + +console = Console() +LOG = logging.getLogger(__name__) + + +def _get_node_roles( + deployment: Deployment, + jhelper: JujuHelper, + console: Console, + show_hints: bool, + nodes: list[str], +) -> Tuple[list[str], list[str], list[str]]: + cluster_status_step: type[ClusterStatusStep] + if deployment.type == "local": + cluster_status_step = LocalClusterStatusStep + else: + cluster_status_step = MaasClusterStatusStep + + results = run_plan([cluster_status_step(deployment, jhelper)], console, show_hints) + cluster_status = get_step_message(results, cluster_status_step) + + control_nodes = [] + compute_nodes = [] + storage_nodes = [] + + for _, node in cluster_status["controller"].items(): + if node["hostname"] in nodes: + if "control" in node["status"]: + control_nodes.append(node["hostname"]) + if "compute" in node["status"]: + compute_nodes.append(node["hostname"]) + if "storage" in node["status"]: + storage_nodes.append(node["hostname"]) + return control_nodes, compute_nodes, storage_nodes + + +@click.command() +@click.option( + "--nodes", + help="Nodes to enable into maintenance mode", + type=click.STRING, + required=True, + callback=validate_nodes, +) +@click.option( + "--force", + help="Force to ignore preflight_checks", + type=click.BOOL, + default=False, +) +@click_option_show_hints +@click.pass_context +def enable_maintenance( + ctx: click.Context, + nodes, + force, + show_hints: bool = False, +) -> None: + console.print("Enable maintenance") + deployment: Deployment = ctx.obj + jhelper = JujuHelper(deployment.get_connected_controller()) + + control_nodes, compute_nodes, _ = _get_node_roles( + deployment=deployment, + jhelper=jhelper, + console=console, + show_hints=show_hints, + nodes=nodes, + ) + + # Run preflight_checks + preflight_checks: list[Check] = [] + # Prefligh checks for control role + preflight_checks += [ + # This check is to avoid issue which maintenance mode haven't support + # control role, which should be removed after control role be supported. + checks.NodeisNotControlRoleCheck(nodes=control_nodes, force=force), + ] + # Prefligh checks for compute role + preflight_checks += [ + checks.InstancesStatusCheck(jhelper=jhelper, nodes=compute_nodes, force=force), + checks.NoEphemeralDiskCheck(jhelper=jhelper, nodes=compute_nodes, force=force), + ] + run_preflight_checks(preflight_checks, console) + + for node in nodes: + plan: list[BaseStep] = [] + post_checks: list[Check] = [] + + if node in compute_nodes: + plan.append(RunWatcherHostMaintenanceStep(jhelper=jhelper, node=node)) + + results = run_plan(plan, console, show_hints) + + # Run post checks + if node in compute_nodes: + # Run post checks only enable maintenance step is completed. + result = get_step_result(results, RunWatcherHostMaintenanceStep) + if result.result_type == ResultType.COMPLETED: + post_checks += [ + checks.NovaInDisableStatusCheck( + jhelper=jhelper, node=node, force=force + ), + checks.NoInstancesOnNodeCheck( + jhelper=jhelper, node=node, force=force + ), + ] + run_preflight_checks(post_checks, console) + console.print("Finish enable maintenance") + + +@click.command() +@click.option( + "--nodes", + help="Nodes to enable into maintenance mode", + type=click.STRING, + required=True, + callback=validate_nodes, +) +@click_option_show_hints +@click.pass_context +def disable_maintenance( + ctx: click.Context, + nodes, + show_hints: bool = False, +) -> None: + deployment: Deployment = ctx.obj + client = deployment.get_client() + jhelper = JujuHelper(deployment.get_connected_controller()) + + _, compute_nodes, _ = _get_node_roles( + deployment=deployment, + jhelper=jhelper, + console=console, + show_hints=show_hints, + nodes=nodes, + ) + + for node in nodes: + plan: list[BaseStep] = [] + if node in compute_nodes: + plan += [ + EnableHypervisorStep( + client=client, + node=node, + jhelper=jhelper, + model=deployment.openstack_machines_model, + ), + RunWatcherWorkloadBalancingStep(jhelper=jhelper), + ] + run_plan(plan, console, show_hints) diff --git a/sunbeam-python/tests/unit/sunbeam/core/test_common.py b/sunbeam-python/tests/unit/sunbeam/core/test_common.py index ce5b6ee0..98be0cf0 100644 --- a/sunbeam-python/tests/unit/sunbeam/core/test_common.py +++ b/sunbeam-python/tests/unit/sunbeam/core/test_common.py @@ -20,7 +20,7 @@ import pytest from sunbeam.clusterd.service import ClusterServiceUnavailableException -from sunbeam.core.common import Role, validate_roles +from sunbeam.core.common import Role, validate_nodes, validate_roles from sunbeam.core.deployment import Deployment @@ -183,3 +183,9 @@ def test_validate_roles(): multiple_comma_separated_roles = ("control,compute", "storage,compute") result = validate_roles(Mock(), Mock(), multiple_comma_separated_roles) assert not set(result) ^ all_roles + + +def test_validate_nodes(): + val = "A,B,C" + result = validate_nodes(Mock(), Mock(), val) + assert sorted(result) == sorted(["A", "B", "C"]) diff --git a/sunbeam-python/tests/unit/sunbeam/provider/maintenance/test_maintenance.py b/sunbeam-python/tests/unit/sunbeam/provider/maintenance/test_maintenance.py new file mode 100644 index 00000000..ea5aa83d --- /dev/null +++ b/sunbeam-python/tests/unit/sunbeam/provider/maintenance/test_maintenance.py @@ -0,0 +1,262 @@ +# Copyright (c) 2024 Canonical Ltd. +# +# 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. +from unittest.mock import Mock, call, patch + +import pytest +from click.testing import CliRunner + +from sunbeam.core.common import ResultType +from sunbeam.provider.maintenance.maintenance import ( + _get_node_roles, + console, + disable_maintenance, + enable_maintenance, +) + + +@pytest.mark.parametrize( + "deployment_type,nodes,status,expected", + [ + ( + "local", + ["node-1", "node-2"], + { + 1: {"hostname": "node-1", "status": ["control"]}, + 2: {"hostname": "node-2", "status": ["control"]}, + }, + (["node-1", "node-2"], [], []), + ), + ( + "maas", + ["node-1"], + { + 1: {"hostname": "node-1", "status": ["control"]}, + }, + (["node-1"], [], []), + ), + ( + "local", + ["node-1", "node-2", "node-3"], + { + 1: {"hostname": "node-1", "status": ["control", "storage"]}, + 2: {"hostname": "node-2", "status": ["control", "compute"]}, + 3: {"hostname": "node-3", "status": ["control", "compute", "storage"]}, + 4: {"hostname": "node-4", "status": ["control", "compute", "storage"]}, + }, + ( + ["node-1", "node-2", "node-3"], + ["node-2", "node-3"], + ["node-1", "node-3"], + ), + ), + ], +) +@patch("sunbeam.provider.maintenance.maintenance.LocalClusterStatusStep") +@patch("sunbeam.provider.maintenance.maintenance.MaasClusterStatusStep") +@patch("sunbeam.provider.maintenance.maintenance.run_plan") +@patch("sunbeam.provider.maintenance.maintenance.get_step_message") +def test_get_node_roles( + mock_get_step_message, + mock_run_plan, + mock_maas_cluster_status_step, + mock_local_cluster_status_step, + deployment_type, + nodes, + status, + expected, +): + mock_deployment = Mock() + mock_jhelper = Mock() + mock_console = Mock() + + mock_deployment.type = deployment_type + if deployment_type == "local": + step = mock_local_cluster_status_step + else: + step = mock_maas_cluster_status_step + + mock_get_step_message.return_value = {"controller": status} + + result = _get_node_roles(mock_deployment, mock_jhelper, mock_console, False, nodes) + + mock_run_plan.assert_called_with( + [step(mock_deployment, mock_jhelper)], mock_console, False + ) + mock_get_step_message.assert_called_with(mock_run_plan.return_value, step) + assert result == expected + + +@patch("sunbeam.provider.maintenance.maintenance.get_step_result") +@patch("sunbeam.provider.maintenance.maintenance.RunWatcherHostMaintenanceStep") +@patch("sunbeam.provider.maintenance.maintenance.run_preflight_checks") +@patch("sunbeam.provider.maintenance.maintenance.run_plan") +@patch("sunbeam.provider.maintenance.maintenance.JujuHelper") +@patch("sunbeam.provider.maintenance.maintenance.checks") +@patch("sunbeam.provider.maintenance.maintenance._get_node_roles") +def test_enable_maintenance( + mock_get_node_roles, + mock_checks, + mock_jhelper, + mock_run_plan, + mock_run_preflight_checks, + mock_run_watcher_host_maintenance_step, + mock_get_step_result, +): + runner = CliRunner() + + compute_nodes = ["node-1", "node-2"] + mock_deployment = Mock() + mock_get_node_roles.return_value = ( + [], + compute_nodes, + [], + ) + mock_get_step_result.return_value.result_type = ResultType.COMPLETED + + steps = [ + mock_run_watcher_host_maintenance_step, + mock_checks.NovaInDisableStatusCheck, + mock_checks.NoInstancesOnNodeCheck, + ] + side_effects = [[Mock() for _ in compute_nodes] for _ in steps] + for step, side_effect in zip(steps, side_effects): + step.side_effect = side_effect + + result = runner.invoke( + enable_maintenance, + ["--nodes", ",".join(compute_nodes)], + obj=mock_deployment, + ) + assert result.exit_code == 0 + + # Verify step are been called + mock_run_plan.assert_has_calls( + call( + [side_effect], + console, + True, + ) + for side_effect in side_effects[0] + ) + mock_run_watcher_host_maintenance_step.assert_has_calls( + [call(jhelper=mock_jhelper.return_value, node=node) for node in compute_nodes], + any_order=True, + ) + + # Verify prefligh checks are been called + preflight_checks_calls = [ + call( + [ + mock_checks.NodeisNotControlRoleCheck.return_value, + mock_checks.InstancesStatusCheck.return_value, + mock_checks.NoEphemeralDiskCheck.return_value, + ], + console, + ), + ] + preflight_checks_calls += [ + call( + list(side_effect), + console, + ) + for side_effect in zip(side_effects[1], side_effects[2]) + ] + mock_run_preflight_checks.assert_has_calls(preflight_checks_calls) + mock_checks.NodeisNotControlRoleCheck.assert_called_with(nodes=[], force=False) + mock_checks.InstancesStatusCheck.assert_called_with( + jhelper=mock_jhelper.return_value, nodes=compute_nodes, force=False + ) + mock_checks.NoEphemeralDiskCheck.assert_called_with( + jhelper=mock_jhelper.return_value, nodes=compute_nodes, force=False + ) + mock_checks.NovaInDisableStatusCheck.assert_has_calls( + [ + call(jhelper=mock_jhelper.return_value, node=node, force=False) + for node in compute_nodes + ], + any_order=True, + ) + mock_checks.NoInstancesOnNodeCheck.assert_has_calls( + [ + call(jhelper=mock_jhelper.return_value, node=node, force=False) + for node in compute_nodes + ], + any_order=True, + ) + + +@patch("sunbeam.provider.maintenance.maintenance.RunWatcherWorkloadBalancingStep") +@patch("sunbeam.provider.maintenance.maintenance.EnableHypervisorStep") +@patch("sunbeam.provider.maintenance.maintenance.JujuHelper") +@patch("sunbeam.provider.maintenance.maintenance.run_plan") +@patch("sunbeam.provider.maintenance.maintenance._get_node_roles") +def test_disable_maintenance( + mock_get_node_roles, + mock_run_plan, + mock_jhelper, + mock_enable_hypervisor_step, + mock_run_watcher_host_maintenance_step, +): + runner = CliRunner() + mock_deployment = Mock() + compute_nodes = ["node-1", "node-2", "node-3"] + + mock_get_node_roles.return_value = ( + [], + compute_nodes, + [], + ) + + steps = [mock_enable_hypervisor_step, mock_run_watcher_host_maintenance_step] + side_effects = [[Mock() for _ in compute_nodes] for _ in steps] + for step, side_effect in zip(steps, side_effects): + step.side_effect = side_effect + + result = runner.invoke( + disable_maintenance, + ["--nodes", ",".join(compute_nodes)], + obj=mock_deployment, + ) + assert result.exit_code == 0 + + mock_run_plan.assert_has_calls( + [ + call( + [ + side_effect[0], # EnableHypervisorStep + side_effect[1], # RunWatcherWorkloadBalancingStep + ], + console, + True, + ) + for side_effect in zip(side_effects[0], side_effects[1]) + ] + ) + mock_enable_hypervisor_step.assert_has_calls( + [ + call( + client=mock_deployment.get_client.return_value, + node=node, + jhelper=mock_jhelper.return_value, + model=mock_deployment.openstack_machines_model, + ) + for node in compute_nodes + ], + any_order=True, + ) + mock_run_watcher_host_maintenance_step.assert_has_calls( + [call(jhelper=mock_jhelper.return_value) for _ in compute_nodes], + any_order=True, + )