bot-bottle/tests/unit/test_orchestrator_host_server.py

"""Unit tests for the host control server (issue #468, chunk 1).

Mostly exercises the pure `dispatch()` (socket-free, like the orchestrator
server tests), plus a real-socket round-trip through `BrokerClient` that proves
the full sign -> POST -> verify -> act seam over HTTP.
"""

from __future__ import annotations

import json
import secrets
import threading
import unittest

from bot_bottle.orchestrator.broker import (
    BrokerAuthError,
    LaunchBroker,
    LaunchRequest,
    StubBroker,
    sign_request,
)
from bot_bottle.orchestrator.broker_client import BrokerClient
from bot_bottle.orchestrator.host_server import (
    broker_secret_from_env,
    dispatch,
    make_host_server,
)


def _body(obj: object) -> bytes:
    return json.dumps(obj).encode()


class _RaisingBroker(LaunchBroker):
    """A broker whose backend launch always fails — exercises the 502 path (an
    operational backend failure, distinct from a fail-closed provenance 401)."""

    def _launch(self, req: LaunchRequest) -> None:
        raise RuntimeError("docker down")

    def _teardown(self, req: LaunchRequest) -> None:
        raise RuntimeError("docker down")


class TestDispatch(unittest.TestCase):
    def setUp(self) -> None:
        self.secret = secrets.token_bytes(16)
        self.broker = StubBroker(self.secret)

    def _token(self, **kwargs: object) -> str:
        return sign_request(LaunchRequest(**kwargs), self.secret)  # type: ignore[arg-type]

    def test_health(self) -> None:
        status, payload = dispatch(self.broker, "GET", "/health", b"")
        self.assertEqual(200, status)
        self.assertEqual("ok", payload["status"])

    def test_broker_launch_verifies_and_acts(self) -> None:
        token = self._token(
            op="launch", bottle_id="b1", source_ip="10.243.0.1",
            image_ref="img", slot=2,
        )
        status, payload = dispatch(self.broker, "POST", "/broker", _body({"token": token}))
        self.assertEqual(200, status)
        self.assertEqual("launch", payload["op"])
        self.assertEqual("b1", payload["bottle_id"])
        self.assertEqual("img", payload["image_ref"])
        self.assertEqual(2, payload["slot"])
        self.assertEqual(["b1"], [r.bottle_id for r in self.broker.launched])

    def test_broker_teardown_acts(self) -> None:
        token = self._token(op="teardown", bottle_id="b1")
        status, _ = dispatch(self.broker, "POST", "/broker", _body({"token": token}))
        self.assertEqual(200, status)
        self.assertEqual(["b1"], [r.bottle_id for r in self.broker.torn_down])

    def test_forged_token_is_401_and_nothing_acted(self) -> None:
        forged = sign_request(
            LaunchRequest(op="launch", bottle_id="b1"), secrets.token_bytes(16))
        status, payload = dispatch(self.broker, "POST", "/broker", _body({"token": forged}))
        self.assertEqual(401, status)
        self.assertIn("broker auth failed", str(payload["error"]))
        self.assertEqual([], self.broker.launched)  # fail-closed: never launched

    def test_backend_failure_is_502(self) -> None:
        broker = _RaisingBroker(self.secret)
        token = self._token(op="launch", bottle_id="b1", image_ref="img")
        status, payload = dispatch(broker, "POST", "/broker", _body({"token": token}))
        self.assertEqual(502, status)
        self.assertIn("backend launch failed", str(payload["error"]))

    def test_missing_token_is_400(self) -> None:
        status, _ = dispatch(self.broker, "POST", "/broker", _body({}))
        self.assertEqual(400, status)

    def test_bad_json_is_400(self) -> None:
        status, _ = dispatch(self.broker, "POST", "/broker", b"{not json")
        self.assertEqual(400, status)

    def test_unknown_route_404(self) -> None:
        status, _ = dispatch(self.broker, "GET", "/nope", b"")
        self.assertEqual(404, status)

    def test_trailing_slash_normalized(self) -> None:
        status, _ = dispatch(self.broker, "GET", "/health/", b"")
        self.assertEqual(200, status)


class TestBrokerSecretFromEnv(unittest.TestCase):
    def test_reads_hex_secret(self) -> None:
        s = secrets.token_bytes(16)
        self.assertEqual(s, broker_secret_from_env({"BOT_BOTTLE_BROKER_SECRET": s.hex()}))

    def test_unset_is_none(self) -> None:
        self.assertIsNone(broker_secret_from_env({}))

    def test_invalid_hex_is_none(self) -> None:
        self.assertIsNone(broker_secret_from_env({"BOT_BOTTLE_BROKER_SECRET": "not-hex"}))


class TestSeamRoundTrip(unittest.TestCase):
    """The whole point of chunk 1: a request signed by the orchestrator side is
    POSTed to a real host control server, verified there, and acted on — over
    HTTP, not an in-process call."""

    def _serve(self, broker: LaunchBroker) -> BrokerClient:
        server = make_host_server(broker, "127.0.0.1", 0)
        self.addCleanup(server.server_close)
        threading.Thread(target=server.serve_forever, daemon=True).start()
        self.addCleanup(server.shutdown)
        host, port = server.server_address[0], server.server_address[1]
        return BrokerClient(f"http://{host}:{port}")

    def test_sign_post_verify_act_over_http(self) -> None:
        secret = secrets.token_bytes(16)
        broker = StubBroker(secret)
        client = self._serve(broker)
        req = LaunchRequest(
            op="launch", bottle_id="b1", source_ip="10.0.0.1", image_ref="img", slot=1)
        got = client.submit(sign_request(req, secret))
        self.assertEqual(req, got)  # the controller echoes the verified request
        self.assertEqual(["b1"], [r.bottle_id for r in broker.launched])

    def test_forged_token_raises_broker_auth_error_over_http(self) -> None:
        secret = secrets.token_bytes(16)
        broker = StubBroker(secret)
        client = self._serve(broker)
        forged = sign_request(
            LaunchRequest(op="launch", bottle_id="b1"), secrets.token_bytes(16))
        with self.assertRaises(BrokerAuthError):
            client.submit(forged)
        self.assertEqual([], broker.launched)  # fail-closed across the wire


if __name__ == "__main__":
    unittest.main()