diff --git a/rebootmgr/main.py b/rebootmgr/main.py index b96a921..52ec70f 100644 --- a/rebootmgr/main.py +++ b/rebootmgr/main.py @@ -351,21 +351,38 @@ def getuser(): return user or getpass.getuser() -def do_set_global_stop_flag(con, dc): - reason = "Set by " + getuser()\ - + " " + str(datetime.datetime.now()) - con.kv.put("service/rebootmgr/stop", reason, dc=dc) - LOG.warning("Set %s global stop flag: %s", dc, reason) +def do_set_global_stop_flag(con, dc, reason=None): + msg_parts = ["Set by", getuser(), str(datetime.datetime.now())] + if reason: + msg_parts.append(reason) + message = " ".join(msg_parts) + con.kv.put("service/rebootmgr/stop", message, dc=dc) + LOG.warning("Set %s global stop flag: %s", dc, message) -def do_set_local_stop_flag(con, hostname): - reason = "Node disabled by " + getuser()\ - + " " + str(datetime.datetime.now()) +def do_unset_global_stop_flag(con, dc): + con.kv.delete("service/rebootmgr/stop", dc=dc) + LOG.warning("Remove %s global stop flag", dc) + + +def do_set_local_stop_flag(con, hostname, reason=None): + msg_parts = ["Node disabled by", getuser(), str(datetime.datetime.now())] + if reason: + msg_parts.append(reason) + message = " ".join(msg_parts) config = get_config(con, hostname) config["enabled"] = False - config["message"] = reason + config["message"] = message + put_config(con, hostname, config) + LOG.warning("Set %s local stop flag: %s", hostname, message) + + +def do_unset_local_stop_flag(con, hostname): + config = get_config(con, hostname) + config["enabled"] = True + config["message"] = "" put_config(con, hostname, config) - LOG.warning("Set %s local stop flag: %s", hostname, reason) + LOG.warning("Unset %s local stop flag", hostname) @click.command() @@ -388,13 +405,17 @@ def do_set_local_stop_flag(con, hostname): default=os.environ.get("REBOOTMGR_CONSUL_PORT", 8500)) @click.option("--ensure-config", help="If there is no valid configuration in consul, create a default one.", is_flag=True) @click.option("--set-global-stop-flag", metavar="CLUSTER", help="Stop the rebootmgr cluster-wide in the specified cluster") +@click.option("--unset-global-stop-flag", metavar="CLUSTER", help="Remove the cluster-wide stop flag in the specified cluster") @click.option("--set-local-stop-flag", help="Stop the rebootmgr on this node", is_flag=True) +@click.option("--unset-local-stop-flag", help="Remove the stop flag on this node", is_flag=True) +@click.option("--stop-reason", help="Reason to set the stop flag") @click.option("--skip-reboot-in-progress-key", help="Don't set the reboot_in_progress consul key before rebooting", is_flag=True) @click.option("--task-timeout", help="Minutes that rebootmgr waits for each task to finish. Default are 120 minutes", default=120, type=int) @click.version_option() def cli(verbose, consul, consul_port, check_triggers, check_uptime, dryrun, maintenance_reason, ignore_global_stop_flag, ignore_node_disabled, ignore_failed_checks, check_holidays, post_reboot_wait_until_healthy, lazy_consul_checks, - ensure_config, set_global_stop_flag, set_local_stop_flag, skip_reboot_in_progress_key, task_timeout): + ensure_config, set_global_stop_flag, unset_global_stop_flag, set_local_stop_flag, unset_local_stop_flag, stop_reason, + skip_reboot_in_progress_key, task_timeout): """Reboot Manager Default values of parameteres are environment variables (if set) @@ -414,11 +435,19 @@ def cli(verbose, consul, consul_port, check_triggers, check_uptime, dryrun, main sys.exit(0) if set_global_stop_flag: - do_set_global_stop_flag(con, set_global_stop_flag) + do_set_global_stop_flag(con, set_global_stop_flag, reason=stop_reason) + sys.exit(0) + + if unset_global_stop_flag: + do_unset_global_stop_flag(con, unset_global_stop_flag) sys.exit(0) if set_local_stop_flag: - do_set_local_stop_flag(con, hostname) + do_set_local_stop_flag(con, hostname, reason=stop_reason) + sys.exit(0) + + if unset_local_stop_flag: + do_unset_local_stop_flag(con, hostname) sys.exit(0) if not config_is_present_and_valid(con, hostname): diff --git a/tests/test_stopflag.py b/tests/test_stopflag.py index db895c5..d8a201c 100644 --- a/tests/test_stopflag.py +++ b/tests/test_stopflag.py @@ -42,7 +42,52 @@ def test_set_global_stop_flag( assert "Set "+datacenter+" global stop flag:" in result.output idx, data = consul_cluster[0].kv.get("service/rebootmgr/stop", dc=datacenter) assert idx is not None - assert data != "" + assert data["Value"] + assert result.exit_code == 0 + + +def test_set_global_stop_flag_with_reason( + run_cli, forward_consul_port, consul_cluster, + mock_subprocess_run, mocker): + mocked_sleep = mocker.patch("time.sleep") + mocked_run = mock_subprocess_run(["shutdown", "-r", "+1"]) + datacenter = "test" + + result = run_cli(rebootmgr, ["-v", "--set-global-stop-flag", datacenter, + "--stop-reason", "My reason"]) + + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + assert "Set "+datacenter+" global stop flag:" in result.output + idx, data = consul_cluster[0].kv.get("service/rebootmgr/stop", dc=datacenter) + assert idx is not None + assert "My reason" in data["Value"].decode() + assert result.exit_code == 0 + + +def test_unset_global_stop_flag( + run_cli, forward_consul_port, consul_cluster, + mock_subprocess_run, mocker): + mocked_sleep = mocker.patch("time.sleep") + mocked_run = mock_subprocess_run(["shutdown", "-r", "+1"]) + datacenter = "test" + + result = run_cli(rebootmgr, ["-v", "--set-global-stop-flag", datacenter]) + + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + idx, data = consul_cluster[0].kv.get("service/rebootmgr/stop", dc=datacenter) + assert idx is not None + assert data + assert result.exit_code == 0 + + result = run_cli(rebootmgr, ["-v", "--unset-global-stop-flag", datacenter]) + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + idx, data = consul_cluster[0].kv.get("service/rebootmgr/stop", dc=datacenter) + assert "Remove "+datacenter+" global stop flag" in result.output + assert idx is not None + assert data is None assert result.exit_code == 0 @@ -64,3 +109,55 @@ def test_set_local_stop_flag( config = json.loads(data["Value"].decode()) assert config['enabled'] is False assert result.exit_code == 0 + + +def test_set_local_stop_flag_with_reason( + run_cli, forward_consul_port, consul_cluster, + mock_subprocess_run, mocker): + mocked_sleep = mocker.patch("time.sleep") + mocked_run = mock_subprocess_run(["shutdown", "-r", "+1"]) + hostname = socket.gethostname().split(".")[0] + + result = run_cli(rebootmgr, ["-v", "--set-local-stop-flag", + "--stop-reason", "My reason"]) + + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + assert "Set "+hostname+" local stop flag:" in result.output + idx, data = consul_cluster[0].kv.get("service/rebootmgr/nodes/{}/config".format( + hostname)) + assert idx is not None + config = json.loads(data["Value"].decode()) + assert config['enabled'] is False + assert "My reason" in config["message"] + assert result.exit_code == 0 + + +def test_unset_local_stop_flag( + run_cli, forward_consul_port, consul_cluster, + mock_subprocess_run, mocker): + mocked_sleep = mocker.patch("time.sleep") + mocked_run = mock_subprocess_run(["shutdown", "-r", "+1"]) + hostname = socket.gethostname().split(".")[0] + + result = run_cli(rebootmgr, ["-v", "--set-local-stop-flag"]) + + consul_key = "service/rebootmgr/nodes/{}/config".format(hostname) + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + idx, data = consul_cluster[0].kv.get(consul_key) + assert idx is not None + config = json.loads(data["Value"].decode()) + assert config['enabled'] is False + assert result.exit_code == 0 + + result = run_cli(rebootmgr, ["-v", "--unset-local-stop-flag"]) + + mocked_sleep.assert_not_called() + mocked_run.assert_not_called() + assert "Unset "+hostname+" local stop flag" in result.output + idx, data = consul_cluster[0].kv.get(consul_key) + assert idx is not None + config = json.loads(data["Value"].decode()) + assert config['enabled'] is True + assert result.exit_code == 0