fixup! Convert branch state management to jobs

Author: mxsrc

src/api/organization/project/branch/control_tasks.py

@@ -34,10 +34,10 @@
     "stop": BranchServiceStatus.STOPPING,
 }
 
-_CONTROL_TRANSITION_FINAL: dict[str, BranchServiceStatus | None] = {
+_CONTROL_TRANSITION_FINAL: dict[str, BranchServiceStatus] = {
     "pause": BranchServiceStatus.PAUSED,
     "resume": BranchServiceStatus.STARTING,
-    "start": None,
+    "start": BranchServiceStatus.STARTING,
     "stop": BranchServiceStatus.STOPPED,
 }
 
@@ -94,8 +94,7 @@ async def _async_perform_control(branch_id: str, action: str) -> dict:
         if branch is None:
             logger.error("Branch %s not found after control action %s", branch_id, action)
             return {"action": action}
-        if final_status is not None:
-            branch.set_status(final_status)
+        branch.set_status(final_status)
         branch.control_task_id = None
         await session.commit()
 

tests/branches/test_basic.py

@@ -144,7 +144,7 @@ def test_branch_resize(client, org, project, branch_id):
     wait_for_status(
         client,
         f"organizations/{org}/projects/{project}/branches/{branch_id}/",
-        "ACTIVE_HEALTHY",
+        ["RESIZING", "ACTIVE_HEALTHY"],
         BRANCH_TIMEOUT_SEC,
     )
     r = client.get(f"organizations/{org}/projects/{project}/branches/{branch_id}/")

tests/branches/test_start_stop.py

@@ -26,7 +26,7 @@ def test_branch_stop(client, org, project, branch_id):
     wait_for_status(
         client,
         f"organizations/{org}/projects/{project}/branches/{branch_id}/",
-        "STOPPED",
+        ["ACTIVE_HEALTHY", "STOPPING", "STOPPED"],
         BRANCH_TIMEOUT_SEC,
     )
 
@@ -38,6 +38,6 @@ def test_branch_start(client, org, project, branch_id):
     wait_for_status(
         client,
         f"organizations/{org}/projects/{project}/branches/{branch_id}/",
-        "ACTIVE_HEALTHY",
+        ["STOPPED", "STARTING", "ACTIVE_HEALTHY"],
         BRANCH_TIMEOUT_SEC,
     )

tests/conftest.py

@@ -14,7 +14,7 @@
 VELA_KEYCLOAK_CLIENT_SECRET = os.getenv("VELA_KEYCLOAK_CLIENT_SECRET", "controller-secret")
 BRANCH_TIMEOUT_SEC = int(os.getenv("BRANCH_TIMEOUT_SEC", "900"))
 
-_POLL_INTERVAL = 10
+_POLL_INTERVAL = 1
 
 
 class _KeycloakAuth(httpx.Auth):
@@ -51,20 +51,41 @@ def auth_flow(self, request: httpx.Request) -> Generator[httpx.Request, httpx.Re
         yield request
 
 
-def wait_for_status(client: httpx.Client, url: str, expected: str, timeout: int = BRANCH_TIMEOUT_SEC) -> dict:
-    """Poll GET url every 10s until response.json()["status"] == expected or timeout."""
+def wait_for_status(
+    client: httpx.Client,
+    url: str,
+    expected: list[str],
+    timeout: int = BRANCH_TIMEOUT_SEC,
+) -> dict:
+    """Poll GET url every 10s until the resource progresses through ``expected``.
+
+    ``expected`` is an ordered list of statuses.  Each status observed must
+    appear at or after the current position in the list; going backward raises
+    immediately.  Statuses not in the list are treated as unknown intermediate
+    states and ignored.  ERROR always raises immediately.  The final entry in
+    ``expected`` is the success condition.
+    """
     deadline = time.monotonic() + timeout
+    idx = 0
+    final = expected[-1]
     while True:
         r = client.get(url, timeout=30)
         r.raise_for_status()
         data = r.json()
         current = data.get("status")
-        if current == expected:
-            return data
         if current == "ERROR":
-            raise RuntimeError(f"Resource entered ERROR state while waiting for status={expected!r}")
+            raise RuntimeError(f"Resource entered ERROR state while waiting for status={final!r}")
+        try:
+            idx = expected.index(current, idx)
+        except ValueError:
+            if current in expected[:idx]:
+                raise RuntimeError(f"Status {current!r} regressed; expected one of {expected[idx:]}") from None
+            # Unknown intermediate state; keep waiting
+        else:
+            if current == final:
+                return data
         if time.monotonic() >= deadline:
-            raise TimeoutError(f"Timed out waiting for status={expected!r}; last status={current!r}")
+            raise TimeoutError(f"Timed out waiting for status={final!r}; last status={current!r}")
         time.sleep(_POLL_INTERVAL)
 
 
@@ -152,7 +173,7 @@ def _factory(org_id: ULID, project_id: ULID, name: str, **overrides: object) ->
         wait_for_status(
             client,
             f"organizations/{org_id}/projects/{project_id}/branches/{uid}/",
-            "ACTIVE_HEALTHY",
+            ["ACTIVE_HEALTHY"],
             BRANCH_TIMEOUT_SEC,
         )
         created.append((org_id, project_id, uid))