Update automation pause and automation resume to handle automations with same name (#13131)

Author: serinamarie

src/prefect/events/cli/automations.py

@@ -17,6 +17,7 @@
 from prefect.cli._utilities import exit_with_error, exit_with_success
 from prefect.cli.root import app
 from prefect.client.orchestration import get_client
+from prefect.events.schemas.automations import Automation
 from prefect.exceptions import PrefectHTTPStatusError
 
 automations_app = PrefectTyper(
@@ -149,7 +150,7 @@ async def inspect(
     if yaml or json:
         if isinstance(automation, list):
             automation = [a.dict(json_compatible=True) for a in automation]
-        elif isinstance(automation, dict):
+        elif isinstance(automation, Automation):
             automation = automation.dict(json_compatible=True)
         if yaml:
             app.console.print(pyyaml.dump(automation, sort_keys=False))
@@ -163,32 +164,112 @@ async def inspect(
 
 @automations_app.command(aliases=["enable"])
 @requires_automations
-async def resume(id_or_name: str):
-    """Resume an automation."""
-    async with get_client() as client:
-        automation = await client.find_automation(id_or_name)
-        if not automation:
-            exit_with_error(f"Automation {id_or_name!r} not found.")
+async def resume(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+):
+    """
+    Resume an automation.
 
-    async with get_client() as client:
-        await client.resume_automation(automation.id)
+    Arguments:
+
+            name: the name of the automation to resume
+
+            id: the id of the automation to resume
+
+    Examples:
+
+            $ prefect automation resume "my-automation"
+
+            $ prefect automation resume --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+    """
+    if not id and not name:
+        exit_with_error("Please provide either a name or an id.")
 
-    exit_with_success(f"Resumed automation {automation.name!r} ({automation.id})")
+    if name:
+        async with get_client() as client:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(
+                    f"Automation with name {name!r} not found. You can also specify an id with the `--id` flag."
+                )
+            if len(automation) > 1:
+                if not typer.confirm(
+                    f"Multiple automations found with name {name!r}. Do you want to resume all of them?",
+                    default=False,
+                ):
+                    exit_with_error("Resume aborted.")
+
+            for a in automation:
+                await client.resume_automation(a.id)
+            exit_with_success(
+                f"Resumed automation(s) with name {name!r} and id(s) {', '.join([repr(str(a.id)) for a in automation])}."
+            )
+
+    elif id:
+        async with get_client() as client:
+            try:
+                uuid_id = UUID(id)
+                automation = await client.read_automation(uuid_id)
+            except (PrefectHTTPStatusError, ValueError):
+                exit_with_error(f"Automation with id {id!r} not found.")
+            await client.resume_automation(automation.id)
+            exit_with_success(f"Resumed automation with id {str(automation.id)!r}.")
 
 
 @automations_app.command(aliases=["disable"])
 @requires_automations
-async def pause(id_or_name: str):
-    """Pause an automation."""
-    async with get_client() as client:
-        automation = await client.find_automation(id_or_name)
-        if not automation:
-            exit_with_error(f"Automation {id_or_name!r} not found.")
+async def pause(
+    name: Optional[str] = typer.Argument(None, help="An automation's name"),
+    id: Optional[str] = typer.Option(None, "--id", help="An automation's id"),
+):
+    """
+    Pause an automation.
 
-    async with get_client() as client:
-        await client.pause_automation(automation.id)
+    Arguments:
+
+            name: the name of the automation to pause
+
+            id: the id of the automation to pause
+
+    Examples:
+
+        $ prefect automation pause "my-automation"
+
+        $ prefect automation pause --id "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
+    """
+    if not id and not name:
+        exit_with_error("Please provide either a name or an id.")
 
-    exit_with_success(f"Paused automation {automation.name!r} ({automation.id})")
+    if name:
+        async with get_client() as client:
+            automation = await client.read_automations_by_name(name=name)
+            if not automation:
+                exit_with_error(
+                    f"Automation with name {name!r} not found. You can also specify an id with the `--id` flag."
+                )
+            if len(automation) > 1:
+                if not typer.confirm(
+                    f"Multiple automations found with name {name!r}. Do you want to pause all of them?",
+                    default=False,
+                ):
+                    exit_with_error("Pause aborted.")
+
+            for a in automation:
+                await client.pause_automation(a.id)
+            exit_with_success(
+                f"Paused automation(s) with name {name!r} and id(s) {', '.join([repr(str(a.id)) for a in automation])}."
+            )
+
+    elif id:
+        async with get_client() as client:
+            try:
+                uuid_id = UUID(id)
+                automation = await client.read_automation(uuid_id)
+            except (PrefectHTTPStatusError, ValueError):
+                exit_with_error(f"Automation with id {id!r} not found.")
+            await client.pause_automation(automation.id)
+            exit_with_success(f"Paused automation with id {str(automation.id)!r}.")
 
 
 @automations_app.command()

tests/events/client/cli/test_automations.py

@@ -211,7 +211,7 @@ def test_inspecting_by_name_not_found(
     )
 
 
-def test_inspecting_in_json(
+def test_inspecting_by_name_in_json(
     various_automations: List[Automation], read_automations_by_name: mock.AsyncMock
 ):
     read_automations_by_name.return_value = [various_automations[0]]
@@ -223,7 +223,26 @@ def test_inspecting_in_json(
     assert loaded[0]["id"] == "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
 
 
-def test_inspecting_in_yaml(
+def test_inspecting_by_id_in_json(
+    various_automations: List[Automation], read_automation: mock.AsyncMock
+):
+    read_automation.return_value = various_automations[1]
+    result = invoke_and_assert(
+        [
+            "automations",
+            "inspect",
+            "--id",
+            "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
+            "--json",
+        ],
+        expected_code=0,
+    )
+    loaded = orjson.loads(result.output)
+    assert loaded["name"] == "My Other Reactive Automation"
+    assert loaded["id"] == "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
+
+
+def test_inspecting_by_name_in_yaml(
     various_automations: List[Automation], read_automations_by_name: mock.AsyncMock
 ):
     read_automations_by_name.return_value = [various_automations[0]]
@@ -235,6 +254,25 @@ def test_inspecting_in_yaml(
     assert loaded[0]["id"] == "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
 
 
+def test_inspecting_by_id_in_yaml(
+    various_automations: List[Automation], read_automation: mock.AsyncMock
+):
+    read_automation.return_value = various_automations[1]
+    result = invoke_and_assert(
+        [
+            "automations",
+            "inspect",
+            "--id",
+            "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
+            "--yaml",
+        ],
+        expected_code=0,
+    )
+    loaded = yaml.safe_load(result.output)
+    assert loaded["name"] == "My Other Reactive Automation"
+    assert loaded["id"] == "bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"
+
+
 @pytest.fixture
 def pause_automation() -> Generator[mock.AsyncMock, None, None]:
     with mock.patch(
@@ -244,26 +282,69 @@ def pause_automation() -> Generator[mock.AsyncMock, None, None]:
 
 
 def test_pausing_by_name(
-    pause_automation: mock.AsyncMock, various_automations: List[Automation]
+    pause_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automations_by_name: mock.AsyncMock,
 ):
+    read_automations_by_name.return_value = [various_automations[0]]
     invoke_and_assert(
         ["automations", "pause", "My First Reactive"],
         expected_code=0,
-        expected_output_contains=["Paused automation 'My First Reactive'"],
+        expected_output_contains=[
+            "Paused automation(s) with name 'My First Reactive' and id(s) 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'"
+        ],
     )
 
     pause_automation.assert_awaited_once_with(
         mock.ANY, UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
     )
 
 
-def test_pausing_not_found(
-    pause_automation: mock.AsyncMock, various_automations: List[Automation]
+def test_pausing_by_name_not_found(
+    pause_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automations_by_name: mock.AsyncMock,
 ):
+    read_automations_by_name.return_value = None
     invoke_and_assert(
         ["automations", "pause", "Wha?"],
         expected_code=1,
-        expected_output_contains=["Automation 'Wha?' not found"],
+        expected_output_contains=["Automation with name 'Wha?' not found"],
+    )
+
+    pause_automation.assert_not_awaited()
+
+
+def test_pausing_by_id(
+    pause_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automation: mock.AsyncMock,
+):
+    read_automation.return_value = various_automations[0]
+    invoke_and_assert(
+        ["automations", "pause", "--id", "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"],
+        expected_code=0,
+        expected_output_contains=[
+            "Paused automation with id 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'"
+        ],
+    )
+
+    pause_automation.assert_awaited_once_with(
+        mock.ANY, UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
+    )
+
+
+def test_pausing_by_id_not_found(
+    pause_automation: mock.AsyncMock,
+    read_automation: mock.AsyncMock,
+):
+    read_automation.return_value = None
+    invoke_and_assert(
+        ["automations", "pause", "--id", "zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz"],
+        expected_code=1,
+        expected_output_contains=[
+            "Automation with id 'zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz' not found"
+        ],
     )
 
     pause_automation.assert_not_awaited()
@@ -278,26 +359,68 @@ def resume_automation() -> Generator[mock.AsyncMock, None, None]:
 
 
 def test_resuming_by_name(
-    resume_automation: mock.AsyncMock, various_automations: List[Automation]
+    resume_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automations_by_name: mock.AsyncMock,
 ):
+    read_automations_by_name.return_value = [various_automations[0]]
     invoke_and_assert(
         ["automations", "resume", "My First Reactive"],
         expected_code=0,
-        expected_output_contains=["Resumed automation 'My First Reactive'"],
+        expected_output_contains=[
+            "Resumed automation(s) with name 'My First Reactive' and id(s) 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'"
+        ],
     )
 
     resume_automation.assert_awaited_once_with(
         mock.ANY, UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
     )
 
 
-def test_resuming_not_found(
-    resume_automation: mock.AsyncMock, various_automations: List[Automation]
+def test_resuming_by_name_not_found(
+    resume_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automations_by_name: mock.AsyncMock,
 ):
+    read_automations_by_name.return_value = None
     invoke_and_assert(
         ["automations", "resume", "Wha?"],
         expected_code=1,
-        expected_output_contains=["Automation 'Wha?' not found"],
+        expected_output_contains=["Automation with name 'Wha?' not found"],
+    )
+
+    resume_automation.assert_not_awaited()
+
+
+def test_resuming_by_id(
+    resume_automation: mock.AsyncMock,
+    various_automations: List[Automation],
+    read_automation: mock.AsyncMock,
+):
+    read_automation.return_value = various_automations[0]
+    invoke_and_assert(
+        ["automations", "resume", "--id", "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"],
+        expected_code=0,
+        expected_output_contains=[
+            "Resumed automation with id 'aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'"
+        ],
+    )
+
+    resume_automation.assert_awaited_once_with(
+        mock.ANY, UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
+    )
+
+
+def test_resuming_by_id_not_found(
+    resume_automation: mock.AsyncMock, read_automation: mock.AsyncMock
+):
+    read_automation.return_value = None
+    invoke_and_assert(
+        ["automations", "resume", "--id", "zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz"],
+        expected_code=1,
+        expected_output_contains=[
+            "Automation with id 'zzzzzzzz-zzzz-zzzz-zzzz-zzzzzzzzzzzz' not found"
+        ],
     )
 
     resume_automation.assert_not_awaited()