DataOps-platform/tests/integration/test_v53_mcp_kestra_runner_l3.py

"""Opt-in V53 MCP -> Kestra -> Runner production-wiring acceptance."""

import json
import os
import time
from pathlib import Path

import anyio
import pytest
import requests
from mcp import ClientSession
from mcp.client.stdio import StdioServerParameters, stdio_client
from neo4j import GraphDatabase
from sqlalchemy import create_engine, text

from app.core.common.identifiers import new_governance_uid
from app.core.data_source.credentials import (
    CredentialCodec,
    DataSourceCredentialRepository,
)
from app.core.data_source.definitions import DataSourceDefinitionRepository
from app.core.data_source.models import (
    DataSourceCredential,
    DataSourceDefinition,
)
from app.core.mcp.bootstrap import build_canary_verifier

pytestmark = pytest.mark.integration
ROOT = Path(__file__).resolve().parents[2]
DATABASE_URL = "postgresql://dataops:dataops-test-password@127.0.0.1:15432/dataops"
KESTRA_URL = "http://127.0.0.1:18080/api/v1"
KESTRA_AUTH = ("admin@dataops.local", "DataOpsKestra1!")
RUNNER_SECRET = "dataops-local-runner-task-token-secret-change-me"


def _workflow_spec(dataflow_uid, source_uid, version):
    return {
        "schema_version": "1.0",
        "dataflow_uid": dataflow_uid,
        "name": f"V53 MCP L3 acceptance v{version}",
        "nodes": [
            {
                "id": "read_customers",
                "type": "sql.query",
                "data_source_uid": source_uid,
                "purpose": "read",
                "config": {
                    "statement": (
                        "SELECT customer_name FROM acceptance_customers "
                        "WHERE id >= :minimum_id ORDER BY id"
                    ),
                    "parameters": {"minimum_id": "${parameters.minimum_id}"},
                },
            }
        ],
        "edges": [],
        "parameters": {"minimum_id": {"type": "integer", "required": True}},
    }


def _schedule_plan():
    return {
        "schema_version": "1.0",
        "timezone": "Asia/Shanghai",
        "triggers": [{"type": "manual"}],
        "max_concurrency": 1,
        "conflict_policy": "skip",
        "timeout_seconds": 120,
        "retry": {"max_attempts": 1, "delay_seconds": 0},
        "backfill": {"max_days": 1, "max_runs": 2},
    }


def _mcp_environment(correlation_id):
    return {
        **os.environ,
        "PYTHONPATH": str(ROOT),
        "DATAOPS_MCP_SUBJECT": "v53-l3-scheduler",
        "DATAOPS_MCP_ROLES": "scheduler",
        "DATAOPS_MCP_BUSINESS_DOMAINS": "sales",
        "DATAOPS_MCP_ENVIRONMENTS": "test",
        "DATAOPS_MCP_CORRELATION_ID": correlation_id,
        "DATAOPS_MCP_DATABASE_URL": DATABASE_URL,
        "KESTRA_BASE_URL": KESTRA_URL,
        "KESTRA_USERNAME": KESTRA_AUTH[0],
        "KESTRA_PASSWORD": KESTRA_AUTH[1],
        "KESTRA_TENANT_ID": "main",
        "RUNNER_TASK_TOKEN_SECRET": RUNNER_SECRET,
        "RUNNER_TASK_TOKEN_TTL_SECONDS": "120",
    }


async def _call(session, name, arguments):
    response = await session.call_tool(name, arguments)
    assert not response.isError, response.content
    if response.structuredContent is not None:
        return response.structuredContent.get("result", response.structuredContent)
    assert len(response.content) == 1
    return json.loads(response.content[0].text)


def _wait_for_terminal(verifier, candidate_id, timeout=30):
    deadline = time.monotonic() + timeout
    evidence = None
    while time.monotonic() < deadline:
        evidence = verifier.verify(candidate_id)
        if evidence["status"] in {"passed", "failed"}:
            return evidence
        time.sleep(0.25)
    raise AssertionError(f"Canary did not finish: {evidence}")


def _flow_disabled(namespace, flow_id):
    response = requests.get(
        f"{KESTRA_URL}/main/flows/{namespace}/{flow_id}",
        auth=KESTRA_AUTH,
        timeout=20,
    )
    response.raise_for_status()
    return bool(response.json().get("disabled"))


def test_v53_real_mcp_candidate_canary_promote_and_rollback():
    if os.environ.get("RUN_V53_MCP_L3") != "1":
        pytest.skip("set RUN_V53_MCP_L3=1 for the V53 L3 acceptance")

    platform_engine = create_engine(DATABASE_URL, pool_pre_ping=True)
    graph_driver = GraphDatabase.driver(
        "bolt://127.0.0.1:17687",
        auth=("neo4j", "Passw0rd"),
    )
    source_uid = new_governance_uid()
    dataflow_uid = new_governance_uid()
    correlation_id = new_governance_uid()
    candidate_ids = []
    deployments = []
    codec = CredentialCodec.from_base64(
        "MDEyMzQ1Njc4OWFiY2RlZjAxMjM0NTY3ODlhYmNkZWY=",
        "v1",
    )

    async def exercise():
        parameters = StdioServerParameters(
            command=str(ROOT / ".venv/bin/python"),
            args=[
                str(ROOT / "mcp-servers/dataops_mcp_stdio.py"),
                "--kind",
                "scheduling",
                "--factory",
                "app.core.mcp.bootstrap:build_scheduling_gateway",
            ],
            env=_mcp_environment(correlation_id),
            cwd=str(ROOT),
        )
        async with (
            stdio_client(parameters) as (read_stream, write_stream),
            ClientSession(read_stream, write_stream) as session,
        ):
            await session.initialize()
            verifier = build_canary_verifier()
            for version in (1, 2):
                candidate = await _call(
                    session,
                    "create_candidate_plan",
                    {
                        "business_domain": "sales",
                        "environment": "test",
                        "workflow_spec": _workflow_spec(
                            dataflow_uid, source_uid, version
                        ),
                        "schedule_plan": _schedule_plan(),
                    },
                )
                candidate_ids.append(candidate["candidate_id"])
                deployment = await _call(
                    session,
                    "deploy_disabled_version",
                    {
                        "candidate_id": candidate["candidate_id"],
                        "business_domain": "sales",
                        "environment": "test",
                    },
                )
                deployments.append(deployment)
                canary = await _call(
                    session,
                    "run_canary",
                    {
                        "candidate_id": candidate["candidate_id"],
                        "business_domain": "sales",
                        "environment": "test",
                        "inputs": {"minimum_id": 1},
                    },
                )
                assert canary["status"] == "started"
                evidence = _wait_for_terminal(verifier, candidate["candidate_id"])
                assert evidence["status"] == "passed"
                promoted = await _call(
                    session,
                    "promote_candidate",
                    {
                        "candidate_id": candidate["candidate_id"],
                        "business_domain": "sales",
                        "environment": "test",
                        "idempotency_key": (
                            f"v53-l3-{correlation_id}-promote-{version}"
                        ),
                    },
                )
                assert promoted["status"] == "promoted"

            replay = await _call(
                session,
                "promote_candidate",
                {
                    "candidate_id": candidate_ids[1],
                    "business_domain": "sales",
                    "environment": "test",
                    "idempotency_key": (f"v53-l3-{correlation_id}-promote-2"),
                },
            )
            assert replay["candidate_id"] == candidate_ids[1]
            assert _flow_disabled(
                deployments[0]["namespace"],
                deployments[0]["flow_id"],
            )
            assert not _flow_disabled(
                deployments[1]["namespace"],
                deployments[1]["flow_id"],
            )

            rollback = await _call(
                session,
                "rollback_to_previous_version",
                {
                    "candidate_id": candidate_ids[1],
                    "business_domain": "sales",
                    "environment": "test",
                    "idempotency_key": (f"v53-l3-{correlation_id}-rollback-2"),
                },
            )
            assert rollback["active_candidate_id"] == candidate_ids[0]
            assert not _flow_disabled(
                deployments[0]["namespace"],
                deployments[0]["flow_id"],
            )
            assert _flow_disabled(
                deployments[1]["namespace"],
                deployments[1]["flow_id"],
            )

    try:
        with platform_engine.begin() as connection:
            sealed = DataSourceCredentialRepository(codec).create_version(
                connection,
                data_source_uid=source_uid,
                credential=DataSourceCredential(
                    "source_reader", "source-test-password"
                ),
                actor_uid=None,
            )
        with graph_driver.session() as session:
            DataSourceDefinitionRepository(session).save(
                DataSourceDefinition(
                    uid=source_uid,
                    name_en="v53-mcp-l3-source",
                    database_type="postgresql",
                    host="source-postgres",
                    port=5432,
                    database="acceptance",
                    credential_ref=source_uid,
                    credential_version=sealed.credential_version,
                    pool_size=1,
                    max_overflow=1,
                )
            )

        os.environ.update(_mcp_environment(correlation_id))
        anyio.run(exercise)

        with platform_engine.connect() as connection:
            runner_successes = connection.execute(
                text(
                    "SELECT COUNT(*) FROM public.runner_task_executions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid) "
                    "AND status = 'success'"
                ),
                {"uid": dataflow_uid},
            ).scalar_one()
            operations = (
                connection.execute(
                    text(
                        "SELECT action, status, COUNT(*) AS count "
                        "FROM public.workflow_gateway_operations "
                        "WHERE workflow_version_id = ANY(CAST(:ids AS uuid[])) "
                        "GROUP BY action, status ORDER BY action"
                    ),
                    {"ids": candidate_ids},
                )
                .mappings()
                .all()
            )
        assert runner_successes == 2
        assert [dict(row) for row in operations] == [
            {
                "action": "promote_candidate",
                "status": "succeeded",
                "count": 2,
            },
            {
                "action": "rollback_to_previous_version",
                "status": "succeeded",
                "count": 1,
            },
        ]
    finally:
        for deployment in deployments:
            requests.delete(
                f"{KESTRA_URL}/main/flows/"
                f"{deployment['namespace']}/{deployment['flow_id']}",
                auth=KESTRA_AUTH,
                timeout=20,
            )
        with graph_driver.session() as session:
            DataSourceDefinitionRepository(session).delete(source_uid)
        with platform_engine.begin() as connection:
            if candidate_ids:
                connection.execute(
                    text(
                        "DELETE FROM public.workflow_gateway_operations "
                        "WHERE correlation_id = CAST(:id AS uuid)"
                    ),
                    {"id": correlation_id},
                )
                connection.execute(
                    text(
                        "DELETE FROM public.workflow_plan_audits "
                        "WHERE workflow_version_id = "
                        "ANY(CAST(:ids AS uuid[]))"
                    ),
                    {"ids": candidate_ids},
                )
                connection.execute(
                    text(
                        "DELETE FROM public.dataflow_workflow_versions "
                        "WHERE id = ANY(CAST(:ids AS uuid[]))"
                    ),
                    {"ids": candidate_ids},
                )
            connection.execute(
                text(
                    "DELETE FROM public.runner_task_executions "
                    "WHERE dataflow_uid = CAST(:uid AS uuid)"
                ),
                {"uid": dataflow_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credential_audit_events "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
            connection.execute(
                text(
                    "DELETE FROM public.datasource_credentials "
                    "WHERE data_source_uid = CAST(:uid AS uuid)"
                ),
                {"uid": source_uid},
            )
        graph_driver.close()
        platform_engine.dispose()