dograh/api/tests/test_active_calls.py

"""Unit tests for the per-worker active-call registry (deploy draining).

The registry backs GET /api/v1/health/active-calls, which scripts/rolling_update.sh
(and a k8s preStop hook) polls to wait for live calls to finish before stopping a
worker. The guarantees that matter for draining: register/unregister are
idempotent, and the count only reaches zero when every registered run is gone.
"""

import asyncio

import pytest
from fastapi import FastAPI
from fastapi.testclient import TestClient

from api.routes import main as main_routes
from api.services.pipecat import active_calls
from api.services.pipecat import run_pipeline as run_pipeline_module


def setup_function():
    # Module-level state — start each test from an empty registry.
    active_calls._active_run_ids.clear()


def _make_active_calls_client(
    monkeypatch,
    configured_secret: str | None = "test-dograh-devops-secret",
) -> TestClient:
    monkeypatch.setattr("api.constants.DOGRAH_DEVOPS_SECRET", configured_secret)
    app = FastAPI()
    app.add_api_route(
        "/api/v1/health/active-calls",
        main_routes.active_calls,
        methods=["GET"],
        response_model=main_routes.ActiveCallsResponse,
    )
    return TestClient(app)


def test_starts_empty():
    assert active_calls.active_call_count() == 0


def test_register_counts_distinct_runs():
    active_calls.register_active_call(1)
    active_calls.register_active_call(2)
    assert active_calls.active_call_count() == 2


def test_register_is_idempotent():
    # Registering the same run twice must not double-count, or the count could
    # never drain to zero.
    active_calls.register_active_call(1)
    active_calls.register_active_call(1)
    assert active_calls.active_call_count() == 1


def test_unregister_removes_run():
    active_calls.register_active_call(1)
    active_calls.register_active_call(2)
    active_calls.unregister_active_call(1)
    assert active_calls.active_call_count() == 1


def test_unregister_unknown_run_is_a_noop():
    # discard() semantics: unregistering a run that was never registered (or was
    # already removed) is safe and cannot push the count negative.
    active_calls.unregister_active_call(999)
    assert active_calls.active_call_count() == 0


def test_full_lifecycle_drains_to_zero():
    active_calls.register_active_call(42)
    assert active_calls.active_call_count() == 1
    active_calls.unregister_active_call(42)
    assert active_calls.active_call_count() == 0


@pytest.mark.asyncio
async def test_run_pipeline_counts_call_during_setup(monkeypatch):
    entered_setup = asyncio.Event()
    release_setup = asyncio.Event()

    async def fake_get_workflow_run(*args, **kwargs):
        entered_setup.set()
        await release_setup.wait()
        raise RuntimeError("setup failed")

    monkeypatch.setattr(
        run_pipeline_module.db_client,
        "get_workflow_run",
        fake_get_workflow_run,
    )

    task = asyncio.create_task(
        run_pipeline_module._run_pipeline(
            transport=object(),
            workflow_id=1,
            workflow_run_id=42,
            user_id=7,
        )
    )

    await asyncio.wait_for(entered_setup.wait(), timeout=1.0)
    assert active_calls.active_call_count() == 1

    release_setup.set()
    with pytest.raises(RuntimeError, match="setup failed"):
        await asyncio.wait_for(task, timeout=1.0)
    assert active_calls.active_call_count() == 0


@pytest.mark.asyncio
async def test_webrtc_entrypoint_counts_call_during_setup(monkeypatch):
    entered_setup = asyncio.Event()
    release_setup = asyncio.Event()

    async def fake_get_workflow(*args, **kwargs):
        entered_setup.set()
        await release_setup.wait()
        raise RuntimeError("setup failed")

    monkeypatch.setattr(
        run_pipeline_module.db_client, "get_workflow", fake_get_workflow
    )

    task = asyncio.create_task(
        run_pipeline_module.run_pipeline_smallwebrtc(
            webrtc_connection=object(),
            workflow_id=1,
            workflow_run_id=43,
            user_id=7,
        )
    )

    await asyncio.wait_for(entered_setup.wait(), timeout=1.0)
    assert active_calls.active_call_count() == 1

    release_setup.set()
    with pytest.raises(RuntimeError, match="setup failed"):
        await asyncio.wait_for(task, timeout=1.0)
    assert active_calls.active_call_count() == 0


@pytest.mark.asyncio
async def test_telephony_entrypoint_counts_call_during_setup(monkeypatch):
    entered_setup = asyncio.Event()
    release_setup = asyncio.Event()

    async def fake_get_workflow(*args, **kwargs):
        entered_setup.set()
        await release_setup.wait()
        raise RuntimeError("setup failed")

    monkeypatch.setattr(
        run_pipeline_module.db_client, "get_workflow", fake_get_workflow
    )

    task = asyncio.create_task(
        run_pipeline_module.run_pipeline_telephony(
            websocket=object(),
            provider_name="twilio",
            workflow_id=1,
            workflow_run_id=44,
            user_id=7,
            call_id="call-1",
            transport_kwargs={},
        )
    )

    await asyncio.wait_for(entered_setup.wait(), timeout=1.0)
    assert active_calls.active_call_count() == 1

    release_setup.set()
    with pytest.raises(RuntimeError, match="setup failed"):
        await asyncio.wait_for(task, timeout=1.0)
    assert active_calls.active_call_count() == 0


def test_active_calls_route_requires_configured_secret(monkeypatch):
    client = _make_active_calls_client(monkeypatch, configured_secret=None)

    response = client.get(
        "/api/v1/health/active-calls",
        headers={"X-Dograh-Devops-Secret": "test-dograh-devops-secret"},
    )

    assert response.status_code == 503


def test_active_calls_route_rejects_missing_secret_header(monkeypatch):
    client = _make_active_calls_client(monkeypatch)

    response = client.get("/api/v1/health/active-calls")

    assert response.status_code == 403


def test_active_calls_route_rejects_wrong_secret(monkeypatch):
    client = _make_active_calls_client(monkeypatch)

    response = client.get(
        "/api/v1/health/active-calls",
        headers={"X-Dograh-Devops-Secret": "wrong"},
    )

    assert response.status_code == 403


def test_active_calls_route_returns_count_with_secret(monkeypatch):
    active_calls.register_active_call(42)
    client = _make_active_calls_client(monkeypatch)

    response = client.get(
        "/api/v1/health/active-calls",
        headers={"X-Dograh-Devops-Secret": "test-dograh-devops-secret"},
    )

    assert response.status_code == 200
    assert response.json() == {"active_calls": 1}