From 771b7117cbeaf4278c7ee611d1517b175bb48f42 Mon Sep 17 00:00:00 2001
From: Edan Bainglass <edan.bainglass@gmail.com>
Date: Thu, 30 Jul 2026 13:31:51 +0200
Subject: [PATCH] Correct plurality and add `amount` parameter

---
 aiida_restapi/routers/daemon.py | 40 +++++++++++++++++++++++----------
 tests/test_daemon.py            | 34 +++++++++++++++++++++++-----
 2 files changed, 56 insertions(+), 18 deletions(-)

diff --git a/aiida_restapi/routers/daemon.py b/aiida_restapi/routers/daemon.py
index befd717..b1a2e03 100644
--- a/aiida_restapi/routers/daemon.py
+++ b/aiida_restapi/routers/daemon.py
@@ -8,7 +8,7 @@
 
 from aiida.cmdline.utils.decorators import with_dbenv
 from aiida.engine.daemon.client import DaemonException, get_daemon_client
-from fastapi import APIRouter
+from fastapi import APIRouter, Query
 
 from aiida_restapi.jsonapi.models import errors
 from aiida_restapi.models.daemon import DaemonStatus, DaemonWorker
@@ -41,15 +41,15 @@ async def get_daemon_status() -> DaemonStatus:
 
 
 @read_router.get(
-    '/worker',
+    '/workers',
     response_model=dict[str, DaemonWorker],
     responses={
         500: {'model': errors.DaemonError},
     },
 )
 @with_dbenv()
-async def get_daemon_worker() -> dict[str, dict[str, t.Any]]:
-    """Return daemon worker metadata."""
+async def get_daemon_workers() -> dict[str, dict[str, t.Any]]:
+    """Return daemon workers metadata."""
     client = get_daemon_client()
 
     if not client.is_daemon_running:
@@ -158,16 +158,24 @@ async def _wait_for_num_workers(
     },
 )
 @with_dbenv()
-async def increase_daemon_worker() -> dict[str, t.Any]:
-    """Increase the number of daemon workers by one."""
+async def increase_daemon_workers(
+    amount: t.Annotated[
+        int,
+        Query(
+            description='The number of workers to increase by',
+            ge=1,
+        ),
+    ] = 1,
+) -> dict[str, t.Any]:
+    """Increase the number of daemon workers by a specified amount (default=1)."""
     client = get_daemon_client()
 
     if not client.is_daemon_running:
         raise DaemonException('The daemon is not running.')
 
     initial = client.get_numprocesses()['numprocesses']
-    client.increase_workers(1)
-    num_workers = await _wait_for_num_workers(initial + 1)
+    client.increase_workers(amount)
+    num_workers = await _wait_for_num_workers(initial + amount)
 
     return {
         'running': True,
@@ -183,16 +191,24 @@ async def increase_daemon_worker() -> dict[str, t.Any]:
     },
 )
 @with_dbenv()
-async def decrease_daemon_worker() -> dict[str, t.Any]:
-    """Decrease the number of daemon workers by one."""
+async def decrease_daemon_workers(
+    amount: t.Annotated[
+        int,
+        Query(
+            description='The number of workers to decrease by',
+            ge=1,
+        ),
+    ] = 1,
+) -> dict[str, t.Any]:
+    """Decrease the number of daemon workers by a specified amount (default=1)."""
     client = get_daemon_client()
 
     if not client.is_daemon_running:
         raise DaemonException('The daemon is not running.')
 
     initial = client.get_numprocesses()['numprocesses']
-    client.decrease_workers(1)
-    num_workers = await _wait_for_num_workers(max(initial - 1, 0))
+    client.decrease_workers(amount)
+    num_workers = await _wait_for_num_workers(max(initial - amount, 0))
 
     return {
         'running': True,
diff --git a/tests/test_daemon.py b/tests/test_daemon.py
index 4e62b87..ebc5ad6 100644
--- a/tests/test_daemon.py
+++ b/tests/test_daemon.py
@@ -26,17 +26,17 @@ def test_status_and_start(client: TestClient):
 
 
 @pytest.mark.usefixtures('stopped_daemon_client')
-def test_worker_when_stopped(client: TestClient):
-    """Test ``/daemon/worker`` when the daemon is not running."""
-    response = client.get('/daemon/worker')
+def test_workers_when_stopped(client: TestClient):
+    """Test ``/daemon/workers`` when the daemon is not running."""
+    response = client.get('/daemon/workers')
     assert response.status_code == 200, response.content
     assert response.json() == {}
 
 
 @pytest.mark.usefixtures('started_daemon_client')
-def test_worker_when_running(client: TestClient):
-    """Test ``/daemon/worker`` when the daemon is running."""
-    response = client.get('/daemon/worker')
+def test_workers_when_running(client: TestClient):
+    """Test ``/daemon/workers`` when the daemon is running."""
+    response = client.get('/daemon/workers')
     assert response.status_code == 200, response.content
 
     results = response.json()
@@ -116,3 +116,25 @@ def test_increase_and_decrease_when_running(client: TestClient):
     results = response.json()
     assert results['running'] is True
     assert results['num_workers'] == initial_workers
+
+
+@pytest.mark.usefixtures('started_daemon_client')
+def test_increase_and_decrease_by_amount(client: TestClient):
+    """Test ``/daemon/increase`` and ``/daemon/decrease`` with a specified amount."""
+    response = client.get('/daemon/status')
+    assert response.status_code == 200, response.content
+    initial_workers = response.json()['num_workers']
+
+    response = client.post('/daemon/increase', params={'amount': 3})
+    assert response.status_code == 200, response.content
+
+    results = response.json()
+    assert results['running'] is True
+    assert results['num_workers'] == initial_workers + 3
+
+    response = client.post('/daemon/decrease', params={'amount': 2})
+    assert response.status_code == 200, response.content
+
+    results = response.json()
+    assert results['running'] is True
+    assert results['num_workers'] == initial_workers + 1