Agentswarm/swarm-controller/tests/test_cloud_init.py

"""cloud-init rendering unit tests (no Azure needed).

Covers the two provisioning fixes:
  * private-ACR worker images get an MI-based `docker login` before pull;
  * a fluent-bit sidecar ships VM logs to Log Analytics when configured,
    and is fully absent otherwise.
"""

from __future__ import annotations

from swarm_controller.config import Settings
from swarm_controller.models import DeploymentRecord
from swarm_controller.provisioning.cloud_init import (
    _acr_registry,
    render_cloud_init,
)

_UAMI_RID = (
    "/subscriptions/sub-1/resourcegroups/rg/providers/Microsoft.ManagedIdentity"
    "/userAssignedIdentities/swarm-controller-id"
)
_BASE_SETTINGS = {
    "azure_subscription_id": "sub-1",
    "controller_uami_client_id": "b9fa1c07-9364",
    "controller_uami_resource_id": _UAMI_RID,
    "cosmos_endpoint": "https://heicode.documents.azure.com:443/",
}


def _settings(**overrides) -> Settings:
    return Settings(_env_file=None, **{**_BASE_SETTINGS, **overrides})


def _record() -> DeploymentRecord:
    return DeploymentRecord(
        id="dep-x",
        deployment_id="dep-x",
        tenant_id="tenant-1",
        user_id="user-1",
        vm_name="dep-x",
        max_agents=2,
        model_ref="gpt-5.4",
        billing_secret_ref="azkv://heicode-vault-test.vault.azure.net/secrets/sk",
    )


# ── _acr_registry ────────────────────────────────────────────────────────────


def test_acr_registry_detects_private_acr():
    assert (
        _acr_registry("heicodetest.azurecr.io/agent-swarm-worker@sha256:abc")
        == "heicodetest.azurecr.io"
    )


def test_acr_registry_ignores_public_ghcr():
    assert _acr_registry("ghcr.io/desplega-ai/agent-swarm-worker:latest") is None


# ── ACR login block ────────────────────────────────────────────────────────────


def test_private_acr_image_gets_docker_login_before_pull():
    s = _settings(
        swarm_worker_image="heicodetest.azurecr.io/agent-swarm-worker@sha256:495cf2"
    )
    out = render_cloud_init(s, _record(), "api-key")
    assert "docker login heicodetest.azurecr.io" in out
    assert "oauth2/exchange" in out
    # login must precede the compose pull so the worker image is reachable
    assert out.index("docker login heicodetest.azurecr.io") < out.index(
        "docker compose --env-file"
    )


def test_public_worker_image_has_no_acr_login():
    s = _settings(swarm_worker_image="ghcr.io/desplega-ai/agent-swarm-worker:latest")
    out = render_cloud_init(s, _record(), "api-key")
    assert "docker login" not in out


# ── Log Analytics sidecar ──────────────────────────────────────────────────────


def test_loganalytics_sidecar_present_when_configured():
    s = _settings(loganalytics_workspace_id="159ee835-a534")
    out = render_cloud_init(s, _record(), "api-key")
    # sidecar service + image
    assert "logship:" in out
    assert "fluent/fluent-bit" in out
    # key resolved from KV via MI (same pattern as litestream key)
    assert "secrets/swarm-loganalytics-key" in out
    # workspace id baked into the resolve block; .env passes it through to the container
    assert 'LA_WORKSPACE_ID="159ee835-a534"' in out
    assert "LA_WORKSPACE_ID=${LA_WORKSPACE_ID}" in out
    # records stamped for per-deployment querying
    assert "deployment_id ${DEPLOYMENT_ID}" in out
    assert "vm_name ${VM_NAME}" in out
    assert "VM_NAME=dep-x" in out
    # fluent-bit config written before compose brings the stack up
    assert out.index("fluent-bit/fluent-bit.conf") < out.index(
        "docker compose --env-file"
    )


def test_loganalytics_sidecar_absent_by_default():
    out = render_cloud_init(_settings(), _record(), "api-key")
    assert "fluent-bit" not in out
    assert "logship" not in out
    assert "LA_SHARED_KEY" not in out
    assert "swarm-loganalytics-key" not in out


def test_loganalytics_and_acr_are_independent():
    # private ACR worker + LA on: both blocks present
    s = _settings(
        swarm_worker_image="heicodetest.azurecr.io/agent-swarm-worker@sha256:abc",
        loganalytics_workspace_id="ws-1",
    )
    out = render_cloud_init(s, _record(), "api-key")
    assert "docker login heicodetest.azurecr.io" in out
    assert "logship:" in out


# ── Daytona MCP integration ───────────────────────────────────────────────────


def test_daytona_absent_by_default():
    out = render_cloud_init(_settings(), _record(), "api-key")
    assert "daytona-mcp:" not in out
    assert "daytona MCP server" not in out
    assert "DAYTONA_API_KEY" not in out


def test_daytona_sidecar_and_registration_when_enabled():
    s = _settings(
        daytona_api_url="https://app.daytona.io/api",
        daytona_api_key_secret="swarm-daytona-key",
    )
    out = render_cloud_init(s, _record(), "api-key")
    # compose sidecar
    assert "daytona-mcp:" in out
    assert "heicodetest.azurecr.io/daytona-mcp:heicode-test" in out
    assert "DAYTONA_API_URL=https://app.daytona.io/api" in out
    # key resolved from Key Vault by name, written into .env, never inlined plaintext
    assert "secrets/swarm-daytona-key" in out
    assert "DAYTONA_API_KEY=${DAYTONA_API_KEY}" in out
    # registration + per-agent install via the swarm API
    assert '"name":"daytona"' in out
    assert "http://daytona-mcp:8090/mcp" in out
    assert "/api/mcp-servers/$_DTN_SERVER_ID/install" in out


def test_daytona_requires_both_url_and_secret():
    # only url → disabled
    out = render_cloud_init(_settings(daytona_api_url="https://app.daytona.io/api"), _record(), "k")
    assert "daytona-mcp:" not in out
    # only secret → disabled
    out = render_cloud_init(_settings(daytona_api_key_secret="swarm-daytona-key"), _record(), "k")
    assert "daytona-mcp:" not in out