app/backend/tests/flow/test_service_node.py

"""The service node: a server on a worker, restarted when it ends, its url following."""

import asyncio
import time
from collections.abc import Callable

from fluksio.flow import remote
from fluksio.flow.messages import DType, MessageSpec
from fluksio.flow.nodes import ServiceNode, service
from fluksio.flow.pipeline import Pipeline
from fluksio.flow.remote import RemoteWorker, RemoteWorkerHub


class Agent:
    """The worker's side of the socket: records calls, and ends one on cancel."""

    def __init__(self) -> None:
        self.sent: list[dict] = []
        self.worker: RemoteWorker | None = None

    async def send_json(self, payload: dict) -> None:
        self.sent.append(payload)
        if payload["op"] == "cancel":
            self.end(payload["call_id"], "the service was stopped")

    def end(self, call_id: str, reason: str) -> None:
        assert self.worker is not None
        self.worker.deliver(
            {"call_id": call_id, "ok": False, "error": {"message": reason}}
        )

    def serves(self) -> list[dict]:
        return [frame for frame in self.sent if frame["op"] == "serve"]


async def until(check: Callable[[], bool]) -> None:
    for _ in range(500):
        if check():
            return
        await asyncio.sleep(0.01)
    raise AssertionError("never happened")


def test_a_service_comes_back_and_its_url_follows(monkeypatch):
    monkeypatch.setattr(service, "BACKOFF", (0.01,))

    async def healthy(client, url):
        return None

    monkeypatch.setattr(ServiceNode, "_probe", staticmethod(healthy))

    async def scenario() -> None:
        agent = Agent()
        agent.worker = RemoteWorker(
            name="gpu-worker",
            labels=["gpu"],
            send=agent.send_json,
            loop=asyncio.get_running_loop(),
            host="192.0.2.5",
        )
        hub = RemoteWorkerHub()
        node = ServiceNode(
            provides=[MessageSpec(name="url", dtype=DType.STR)],
            params={"command": "serve --port {port}", "worker": "gpu", "port": 8110},
        )
        node.assign_flow("voice", "parrot")
        pipeline = Pipeline(nodes=[node])
        health: list[str] = []
        node._on_health = lambda _node, status, _detail: health.append(status)
        node.hub = hub

        def url() -> object:
            return pipeline.state.get("voice.url")

        await node.start()
        await asyncio.sleep(0.05)
        # No worker carries the label yet.
        assert health[-1] == "degraded"
        assert not agent.sent

        hub.attach(agent.worker)
        await until(lambda: bool(agent.serves()))
        await until(lambda: url() == "http://192.0.2.5:8110")
        first = agent.serves()[0]
        assert first["argv"] == ["serve", "--port", "8110"]
        assert health[-1] == "ok"

        # The server dies: the url goes at once, and a second call starts.
        agent.end(first["call_id"], "the service exited with code 1")
        await until(lambda: url() == "")
        assert "degraded" in health[health.index("ok") :]
        await until(lambda: len(agent.serves()) == 2)
        await until(lambda: url() == "http://192.0.2.5:8110")

        # Stopping the flow ends the call on the worker, and empties the url.
        await node.stop()
        second = agent.serves()[1]["call_id"]
        assert {"op": "cancel", "call_id": second} in agent.sent
        assert url() == ""

    asyncio.run(scenario())


def test_a_service_cancelled_with_its_loop_ends_its_call(monkeypatch):
    """A loop that cancels every task at once must not strand the call's thread.

    The call's own future is cancelled with the rest, so it reads done while the
    thread still waits for the worker, which then keeps the loop's shutdown
    waiting until the silence bound.
    """
    monkeypatch.setattr(remote, "SILENCE_S", 3.0)

    async def healthy(client, url):
        return None

    monkeypatch.setattr(ServiceNode, "_probe", staticmethod(healthy))
    agent = Agent()

    async def scenario() -> None:
        agent.worker = RemoteWorker(
            name="gpu-worker",
            labels=["gpu"],
            send=agent.send_json,
            loop=asyncio.get_running_loop(),
            host="192.0.2.5",
        )
        hub = RemoteWorkerHub()
        hub.attach(agent.worker)
        node = ServiceNode(
            provides=[MessageSpec(name="url", dtype=DType.STR)],
            params={"command": "serve", "worker": "gpu"},
        )
        node.assign_flow("voice", "parrot")
        pipeline = Pipeline(nodes=[node])
        node.hub = hub
        await node.start()
        await until(lambda: bool(agent.serves()))
        await until(lambda: pipeline.state.get("voice.url") == "http://192.0.2.5:8100")
        # No stop: asyncio.run cancels the node's task and the call together.

    started = time.monotonic()
    asyncio.run(scenario())
    assert time.monotonic() - started < 2
    assert [frame["op"] for frame in agent.sent] == ["serve", "cancel"]


def test_a_service_that_never_gets_healthy_reads_down(monkeypatch):
    monkeypatch.setattr(service, "BACKOFF", (0.01,))

    async def refused(client, url):
        return "is not answering"

    monkeypatch.setattr(ServiceNode, "_probe", staticmethod(refused))

    async def scenario() -> None:
        agent = Agent()
        agent.worker = RemoteWorker(
            name="gpu-worker",
            labels=["gpu"],
            send=agent.send_json,
            loop=asyncio.get_running_loop(),
        )
        hub = RemoteWorkerHub()
        hub.attach(agent.worker)
        node = ServiceNode(params={"command": "serve", "worker": "gpu"})
        node.assign_flow("voice", "parrot")
        health: list[tuple[str, str]] = []
        node._on_health = lambda _node, status, detail: health.append((status, detail))
        node.hub = hub

        await node.start()
        for attempt in range(1, service.DOWN_AFTER + 1):
            await until(lambda n=attempt: len(agent.serves()) == n)
            agent.end(agent.serves()[-1]["call_id"], "Address already in use")
        await until(lambda: health[-1][0] == "down")
        assert "Address already in use" in health[-1][1]
        await node.stop()

    asyncio.run(scenario())