DataOps-platform/tests/integration/test_dataflow_create_reconciliation.py

from __future__ import annotations

import os

import pytest
from neo4j import GraphDatabase
from sqlalchemy import create_engine, text
from sqlalchemy.orm import Session

from app.core.common.identifiers import new_governance_uid
from app.core.data_flow.create_reconciliation import DataFlowCreateReconciler
from app.core.data_flow.dataflows import DataFlowService
from app.core.data_rules.repository import DataRuleRepository


def test_real_postgres_neo4j_orphan_is_reconciled_from_stored_intent(
    monkeypatch,
):
    pg_url = os.environ.get("DATA_RULE_POSTGRES_ACCEPTANCE_URL")
    neo4j_uri = os.environ.get("DATA_RULE_NEO4J_ACCEPTANCE_URI")
    password = os.environ.get("DATA_RULE_NEO4J_ACCEPTANCE_PASSWORD")
    if not pg_url or not neo4j_uri or not password:
        pytest.skip("real PostgreSQL and Neo4j acceptance are not configured")
    user = os.environ.get("DATA_RULE_NEO4J_ACCEPTANCE_USER", "neo4j")
    monkeypatch.setenv("DATABASE_URL", pg_url)
    from app import create_app

    app = create_app()
    app.config.update(
        TESTING=True,
        NEO4J_URI=neo4j_uri,
        NEO4J_USER=user,
        NEO4J_PASSWORD=password,
        NEO4J_ENCRYPTED=False,
    )
    pg = create_engine(pg_url)
    graph = GraphDatabase.driver(
        neo4j_uri, auth=(user, password), encrypted=False
    )
    actor = new_governance_uid()
    reservation_id = None
    dataflow_uid = None
    tag_uid = new_governance_uid()
    tag_node_id = None
    try:
        with graph.session() as graph_session:
            tag_node_id = graph_session.run(
                "CREATE (t:DataLabel {acceptance_uid: $uid}) "
                "RETURN id(t) AS node_id",
                {"uid": tag_uid},
            ).single()["node_id"]
        with app.app_context(), Session(pg) as session:
            session.execute(
                text(
                    "INSERT INTO public.users "
                    "(id, username, display_name, password_hash, status) "
                    "VALUES (CAST(:id AS uuid), :username, "
                    "'Reconcile Acceptance', 'not-a-login-hash', 'active')"
                ),
                {"id": actor, "username": f"reconcile-{actor}"},
            )
            session.commit()
            repository = DataRuleRepository(session)
            receipt = repository.reserve_dataflow_draft(actor_uid=actor)
            reservation_id = receipt["reservation_id"]
            dataflow_uid = receipt["dataflow_uid"]
            session.commit()
            node = {
                "uid": dataflow_uid,
                "name_zh": f"孤儿生产线-{dataflow_uid}",
                "name_en": f"orphan-{dataflow_uid}",
                "category": "应用类",
                "organization": "acceptance",
                "leader": "system",
                "frequency": "月",
                "describe": "reconciliation acceptance",
                "status": "active",
                "update_mode": "append",
                "script_type": "governed",
                "script_requirement": "{}",
                "script_path": "",
            }
            intent = {
                "dataflow_uid": dataflow_uid,
                "node": node,
                "tags": [{"id": tag_node_id}],
            }
            create_request = {
                "payload": {"name_zh": node["name_zh"]},
                "intent": intent,
            }
            closed = {
                key: receipt[key]
                for key in ("reservation_id", "dataflow_uid", "nonce")
            }
            claim = repository.begin_dataflow_create(
                closed,
                actor_uid=actor,
                create_request=create_request,
                create_intent=intent,
                lease_seconds=10,
            )
            session.commit()
            DataFlowService._merge_governed_dataflow(node)
            session.execute(
                text(
                    "UPDATE public.dataflow_draft_reservations "
                    "SET lease_expires_at = CURRENT_TIMESTAMP - INTERVAL '1 second' "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": reservation_id},
            )
            session.commit()
            monkeypatch.setattr(
                DataFlowService,
                "validate_governed_create_intent",
                lambda value, *, repository: value,
            )
            reconciler = DataFlowCreateReconciler(repository)
            before_preview = session.execute(
                text(
                    "SELECT state, attempt_count "
                    "FROM public.dataflow_draft_reservations "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": reservation_id},
            ).one()
            dry_run = reconciler.run(dry_run=True)
            assert dry_run["candidate_count"] == 1
            after_preview = session.execute(
                text(
                    "SELECT state, attempt_count "
                    "FROM public.dataflow_draft_reservations "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": reservation_id},
            ).one()
            assert after_preview == before_preview
            report = reconciler.run(dry_run=False)
            assert report["reconciled_count"] == 1
            assert report["items"][0]["attempt"] == claim["attempt"] + 1
            assert reconciler.run(dry_run=False)["candidate_count"] == 0
            state = session.execute(
                text(
                    "SELECT state, request_digest, result_digest "
                    "FROM public.dataflow_draft_reservations "
                    "WHERE id = CAST(:id AS uuid)"
                ),
                {"id": reservation_id},
            ).mappings().one()
            assert state["state"] == "completed"
            assert state["request_digest"]
            assert state["result_digest"]
            DataFlowService._handle_tag_relationships(
                report["items"][0]["dataflow_node_id"],
                [{"id": tag_node_id}],
                strict=True,
            )
        with graph.session() as graph_session:
            record = graph_session.run(
                "MATCH (n:DataFlow {uid: $uid}) "
                "OPTIONAL MATCH (n)-[r:LABEL]->"
                "(:DataLabel {acceptance_uid: $tag_uid}) "
                "RETURN count(DISTINCT n) AS node_count, "
                "count(r) AS relationship_count",
                {"uid": dataflow_uid, "tag_uid": tag_uid},
            ).single()
        assert record["node_count"] == 1
        assert record["relationship_count"] == 1
    finally:
        if dataflow_uid or tag_uid:
            with graph.session() as graph_session:
                graph_session.run(
                    "MATCH (n) WHERE "
                    "(n:DataFlow AND n.uid = $uid) OR "
                    "(n:DataLabel AND n.acceptance_uid = $tag_uid) "
                    "DETACH DELETE n",
                    {"uid": dataflow_uid, "tag_uid": tag_uid},
                )
        graph.close()
        with Session(pg) as session:
            if reservation_id:
                session.execute(
                    text(
                        "DELETE FROM public.dataflow_draft_reservations "
                        "WHERE id = CAST(:id AS uuid)"
                    ),
                    {"id": reservation_id},
                )
            session.execute(
                text(
                    "DELETE FROM public.users WHERE id = CAST(:id AS uuid)"
                ),
                {"id": actor},
            )
            session.commit()
        pg.dispose()