From 1d3ced52c1b133f76bd50bde6d5cd0fa9a014ccb Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20S=C3=A1nchez-Gallego?= <gallegoj@uw.edu>
Date: Mon, 25 Mar 2024 17:35:43 -0700
Subject: [PATCH] Add actor-restart command

---
 pyproject.toml           |  2 +-
 src/lvmopstools/actor.py | 18 +++++++++++++++++-
 tests/test_actor.py      | 17 +++++++++++++++--
 3 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 4c2291d..e14e1b9 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -28,7 +28,7 @@ include = ["src/lvmopstools/*.yaml"]
 lvmtools = 'lvmopstools.__main__:lvmopstools'
 
 [tool.poetry.dependencies]
-python = "^3.10,<3.13"
+python = "^3.10"
 sdsstools = "^1.3.1"
 sdss-clu = "^2.2.3"
 astropy = "^6.0.0"
diff --git a/src/lvmopstools/actor.py b/src/lvmopstools/actor.py
index 7ab52c0..b88e347 100644
--- a/src/lvmopstools/actor.py
+++ b/src/lvmopstools/actor.py
@@ -168,6 +168,21 @@ async def actor_state(command: Command[LVMActor], *args, **kwargs):
     return command.finish(state={"code": code, "flags": flags, "error": None})
 
 
+@click.command(cls=CluCommand, name="actor-restart")
+@click.option(
+    "--mode",
+    "-m",
+    type=click.Choice(["exit", "reload"]),
+    default="exit",
+    help="How to restart the actor.",
+)
+async def actor_restart(command: Command[LVMActor], mode: str = "exit"):
+    """Restarts the actor."""
+
+    await command.actor.restart(mode=mode)
+    return command.finish()
+
+
 class CheckError(Exception):
     """An exception raised when the :obj:`.LVMActor` check fails."""
 
@@ -208,8 +223,9 @@ def __init__(
         self.restart_after = restart_after
         self.restart_mode = restart_mode
 
-        # Actor state command.
+        # Additional commands.
         self.parser.add_command(actor_state)
+        self.parser.add_command(actor_restart)
 
         # Add keywords in schema for the actor state.
         assert self.model and self.model.schema, "Model schema not defined"
diff --git a/tests/test_actor.py b/tests/test_actor.py
index 7652ffa..04fef42 100644
--- a/tests/test_actor.py
+++ b/tests/test_actor.py
@@ -48,7 +48,7 @@ class NewErrorCodes(enum.Enum):
             SOME_FAILURE_MODE = 1
 
 
-async def test_actor_state(lvm_actor: LVMActor):
+async def test_command_actor_state(lvm_actor: LVMActor):
     assert isinstance(lvm_actor, LVMActor)
 
     cmd = await lvm_actor.invoke_mock_command("actor-state")
@@ -60,7 +60,7 @@ async def test_actor_state(lvm_actor: LVMActor):
     lvm_actor._troubleshoot_internal.assert_not_called()
 
 
-async def test_actor_state_no_model(lvm_actor: LVMActor):
+async def test_command_actor_state_no_model(lvm_actor: LVMActor):
     assert isinstance(lvm_actor, LVMActor)
 
     lvm_actor.model = None
@@ -73,6 +73,19 @@ async def test_actor_state_no_model(lvm_actor: LVMActor):
     assert cmd.replies[-1].body["state"]["error"] is None
 
 
+async def test_command_actor_restart(lvm_actor: LVMActor, mocker: MockerFixture):
+    assert isinstance(lvm_actor, LVMActor)
+
+    lvm_actor.restart = mocker.AsyncMock()
+
+    cmd = await lvm_actor.invoke_mock_command("actor-restart")
+    await cmd
+
+    assert cmd.status.did_succeed
+
+    lvm_actor.restart.assert_called_once()
+
+
 def test_get_error_codes():
     assert ErrorCodes.UNKNOWN == ErrorCodes.get_error_code(9999)