app/backend/tests/flow/test_worker_agent.py

"""The remote worker's own guards: where its token goes, and what node code gets."""

import argparse
import asyncio
import json
import time
from pathlib import Path

import pytest
from fluksio_worker import agent
from fluksio_worker.worker_main import _Reporter


@pytest.mark.parametrize(
    ("url", "trusted"),
    [
        ("wss://engine.example.com/api/v1/workers/attach", True),
        ("https://engine.example.com/api/v1/artifacts", True),
        ("ws://127.0.0.1:8000/api/v1/workers/attach", True),
        ("ws://127.1.2.3:8000/api/v1/workers/attach", True),
        ("ws://[::1]:8000/api/v1/workers/attach", True),
        ("ws://localhost:8000/api/v1/workers/attach", True),
        ("ws://192.0.2.10:8000/api/v1/workers/attach", False),
        ("ws://engine.local:8000/api/v1/workers/attach", False),
        ("http://192.0.2.10:8000/api/v1/artifacts", False),
    ],
)
def test_only_tls_or_this_machine_is_handed_the_token(url: str, trusted: bool) -> None:
    assert agent._encrypted_or_local(url) is trusted


def test_a_cleartext_engine_elsewhere_needs_insecure(
    monkeypatch: pytest.MonkeyPatch, capsys: pytest.CaptureFixture[str]
) -> None:
    url = "ws://192.0.2.10:8000/api/v1/workers/attach"
    with pytest.raises(SystemExit):
        agent.main(["--url", url, "--token", "t"])
    assert "--insecure" in capsys.readouterr().err

    async def attached(self: agent.Agent) -> None:
        return None

    monkeypatch.setattr(agent.Agent, "serve_forever", attached)
    assert agent.main(["--url", url, "--token", "t", "--insecure"]) == 0


def test_node_code_does_not_get_the_hosts_credentials(
    monkeypatch: pytest.MonkeyPatch,
) -> None:
    monkeypatch.setenv("FLUKSIO_WORKER_TOKEN", "worker-token")
    monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "not-for-nodes")
    monkeypatch.setenv("CUDA_VISIBLE_DEVICES", "2")
    args = argparse.Namespace(
        url="wss://engine.example.com/api/v1/workers/attach",
        artifact_url="",
        token="worker-token",
    )

    env = agent.Agent(args).env

    assert "FLUKSIO_WORKER_TOKEN" not in env
    assert "AWS_SECRET_ACCESS_KEY" not in env
    assert env["CUDA_VISIBLE_DEVICES"] == "2"
    assert "PATH" in env
    # Put back on purpose: the artifact store is reached with it.
    assert env["FLUKSIO_ARTIFACT_URL"] == "https://engine.example.com/api/v1/artifacts"


@pytest.mark.parametrize(
    "digest",
    [
        "sha256:../../../etc/passwd",
        "sha256:" + "a" * 63,
        "sha256:" + "A" * 64,
        "md5:" + "a" * 64,
    ],
)
def test_only_a_whole_digest_is_opened(digest: str) -> None:
    """The digest is joined into a path, so anything else could read any file."""
    with pytest.raises(ValueError, match="not an artifact reference"):
        _Reporter("fluksio").load_artifact({"digest": digest})


def _gone(pid: int) -> bool:
    """No such process, or one only waiting to be reaped."""
    try:
        return Path(f"/proc/{pid}/stat").read_text().split()[2] == "Z"
    except FileNotFoundError:
        return True


def test_stopping_a_service_takes_its_process_group_with_it(tmp_path: Path) -> None:
    """A child of the server left running would keep its port."""
    args = argparse.Namespace(
        url="wss://engine.example.com/api/v1/workers/attach",
        artifact_url="",
        token="t",
    )
    worker = agent.Agent(args)
    child = tmp_path / "child"
    sent: list[dict] = []

    class Socket:
        async def send(self, text: str) -> None:
            sent.append(json.loads(text))

    async def scenario() -> None:
        script = f"sleep 300 & echo $! > {child}; wait"
        request = {"op": "serve", "call_id": "svc", "argv": ["sh", "-c", script]}
        task = asyncio.create_task(worker._serve(Socket(), request))
        for _ in range(500):
            if child.exists() and child.read_text().strip():
                break
            await asyncio.sleep(0.01)
        worker._cancel("svc")
        await task

    asyncio.run(scenario())

    pid = int(child.read_text())
    deadline = time.monotonic() + 5
    while not _gone(pid) and time.monotonic() < deadline:
        time.sleep(0.05)
    assert _gone(pid)
    assert not worker.serving
    assert sent == [
        {
            "call_id": "svc",
            "ok": False,
            "error": {
                "type": "ServiceStopped",
                "message": "the service was stopped",
                "short": "the service was stopped",
                "traceback": "",
            },
        }
    ]