Merge pull request #1038 from indrat/procrastinate-shell-onecmd

procrastinate/cli.py

@@ -6,6 +6,7 @@
 import json
 import logging
 import os
+import shlex
 import sys
 from typing import Any, Awaitable, Callable, Literal, Union
 
@@ -475,6 +476,12 @@ def configure_shell_parser(subparsers: argparse._SubParsersAction):
         help="Administration shell for procrastinate",
         **parser_options,
     )
+    add_argument(
+        shell_parser,
+        "args",
+        nargs="*",
+        help="Invoke a shell command and exit",
+    )
     shell_parser.set_defaults(func=shell_)
 
 
@@ -627,15 +634,18 @@ async def healthchecks(app: procrastinate.App):
     print("Found procrastinate_jobs table: OK")
 
 
-async def shell_(app: procrastinate.App):
+async def shell_(app: procrastinate.App, args: list[str]):
     """
     Administration shell for procrastinate.
     """
     shell_obj = shell.ProcrastinateShell(
         job_manager=app.job_manager,
     )
 
-    await utils.sync_to_async(shell_obj.cmdloop)
+    if args:
+        await utils.sync_to_async(shell_obj.onecmd, line=shlex.join(args))
+    else:
+        await utils.sync_to_async(shell_obj.cmdloop)
 
 
 def main():

tests/unit/test_cli.py

@@ -2,6 +2,7 @@
 
 import argparse
 import datetime
+import io
 import json
 import logging
 
@@ -100,6 +101,20 @@ def test_main(mocker):
                 "action": "migrations_path",
             },
         ),
+        (
+            ["shell"],
+            {
+                "command": "shell",
+                "args": [],
+            },
+        ),
+        (
+            ["shell", "list_jobs"],
+            {
+                "command": "shell",
+                "args": ["list_jobs"],
+            },
+        ),
     ],
 )
 def test_parser(input, output):
@@ -253,3 +268,37 @@ def get_sync_connector(self) -> BaseConnector:
     )
     with pytest.raises(argparse.ArgumentError, match="is not async"):
         cli.load_app("foobar")
+
+
+async def test_shell_single_command(app, capsys):
+    @app.task(name="foobar")
+    def mytask(a):
+        pass
+
+    await mytask.defer_async(a=1)
+
+    await cli.shell_(app=app, args=["list_jobs"])
+
+    out, _ = capsys.readouterr()
+
+    assert out == "#1 foobar on default - [todo]\n"
+
+
+async def test_shell_interactive_command(app, capsys, mocker):
+    @app.task(name="foobar")
+    def mytask(a):
+        pass
+
+    await mytask.defer_async(a=1)
+
+    mocker.patch("sys.stdin", io.StringIO("list_jobs\nexit\n"))
+
+    await cli.shell_(app=app, args=[])
+
+    out, _ = capsys.readouterr()
+
+    expected = """Welcome to the procrastinate shell.   Type help or ? to list commands.
+
+procrastinate> #1 foobar on default - [todo]
+procrastinate> """
+    assert out == expected