dograh/api/tests/test_workflow_qa_masking.py

from api.services.configuration.masking import (
    mask_key,
    mask_workflow_definition,
    merge_workflow_api_keys,
)


def _make_workflow_def(nodes):
    """Helper to build a minimal workflow definition dict."""
    return {"nodes": nodes, "edges": [], "viewport": {"x": 0, "y": 0, "zoom": 1}}


def _qa_node(node_id="qa-1", api_key="", **extra_data):
    """Helper to build a QA node."""
    data = {"name": "QA Analysis", "qa_enabled": True, **extra_data}
    if api_key:
        data["qa_api_key"] = api_key
    return {"id": node_id, "type": "qa", "position": {"x": 0, "y": 0}, "data": data}


def _tuner_node(node_id="tuner-1", api_key="", **extra_data):
    """Helper to build a Tuner node."""
    data = {
        "name": "Tuner",
        "tuner_enabled": True,
        "tuner_agent_id": "sales-bot",
        "tuner_workspace_id": 7,
        **extra_data,
    }
    if api_key:
        data["tuner_api_key"] = api_key
    return {
        "id": node_id,
        "type": "tuner",
        "position": {"x": 0, "y": 0},
        "data": data,
    }


def _agent_node(node_id="agent-1"):
    """Helper to build a non-QA node."""
    return {
        "id": node_id,
        "type": "agentNode",
        "position": {"x": 0, "y": 0},
        "data": {"name": "Agent", "prompt": "hello"},
    }


# ---------------------------------------------------------------------------
# mask_workflow_definition
# ---------------------------------------------------------------------------


class TestMaskWorkflowDefinition:
    def test_masks_qa_api_key(self):
        """QA node api_key is masked, showing only last 4 chars."""
        real_key = "sk-proj-abcdefghijklmnop"
        wf = _make_workflow_def([_qa_node(api_key=real_key)])

        masked = mask_workflow_definition(wf)

        masked_key = masked["nodes"][0]["data"]["qa_api_key"]
        assert masked_key == mask_key(real_key)
        assert masked_key.endswith("mnop")
        assert masked_key.startswith("*")
        assert real_key not in str(masked)

    def test_does_not_mutate_original(self):
        """The original workflow definition is not modified."""
        real_key = "sk-proj-abcdefghijklmnop"
        wf = _make_workflow_def([_qa_node(api_key=real_key)])

        mask_workflow_definition(wf)

        assert wf["nodes"][0]["data"]["qa_api_key"] == real_key

    def test_non_qa_nodes_untouched(self):
        """Non-QA nodes are not modified."""
        wf = _make_workflow_def([_agent_node(), _qa_node(api_key="sk-secret1234")])

        masked = mask_workflow_definition(wf)

        assert masked["nodes"][0]["type"] == "agentNode"
        assert "qa_api_key" not in masked["nodes"][0]["data"]
        assert masked["nodes"][1]["data"]["qa_api_key"] == mask_key("sk-secret1234")

    def test_masks_tuner_api_key(self):
        """Tuner node api_key is masked, showing only last 4 chars."""
        real_key = "tuner_live_abcdefghijklmnop"
        wf = _make_workflow_def([_tuner_node(api_key=real_key)])

        masked = mask_workflow_definition(wf)

        masked_key = masked["nodes"][0]["data"]["tuner_api_key"]
        assert masked_key == mask_key(real_key)
        assert masked_key.endswith("mnop")
        assert masked_key.startswith("*")
        assert real_key not in str(masked)

    def test_qa_node_without_api_key(self):
        """QA node with no api_key is left as-is."""
        wf = _make_workflow_def([_qa_node()])

        masked = mask_workflow_definition(wf)

        assert "qa_api_key" not in masked["nodes"][0]["data"]

    def test_qa_node_with_empty_api_key(self):
        """QA node with empty string api_key is left as-is."""
        node = _qa_node()
        node["data"]["qa_api_key"] = ""
        wf = _make_workflow_def([node])

        masked = mask_workflow_definition(wf)

        assert masked["nodes"][0]["data"]["qa_api_key"] == ""

    def test_multiple_qa_nodes(self):
        """All QA nodes in a definition are masked."""
        wf = _make_workflow_def(
            [
                _qa_node(node_id="qa-1", api_key="key-aaaa1111"),
                _qa_node(node_id="qa-2", api_key="key-bbbb2222"),
            ]
        )

        masked = mask_workflow_definition(wf)

        assert masked["nodes"][0]["data"]["qa_api_key"] == mask_key("key-aaaa1111")
        assert masked["nodes"][1]["data"]["qa_api_key"] == mask_key("key-bbbb2222")

    def test_none_definition(self):
        """None input returns None."""
        assert mask_workflow_definition(None) is None

    def test_empty_definition(self):
        """Empty dict returns empty dict."""
        assert mask_workflow_definition({}) == {}

    def test_definition_without_nodes(self):
        """Definition with no nodes key is returned as-is."""
        wf = {"edges": [], "viewport": {"x": 0, "y": 0, "zoom": 1}}
        result = mask_workflow_definition(wf)
        assert result == wf


# ---------------------------------------------------------------------------
# merge_workflow_api_keys
# ---------------------------------------------------------------------------


class TestMergeWorkflowApiKeys:
    def test_masked_key_is_restored(self):
        """When incoming key matches the mask of the existing key, real key is preserved."""
        real_key = "sk-proj-abcdefghijklmnop"
        masked_val = mask_key(real_key)

        existing = _make_workflow_def([_qa_node(api_key=real_key)])
        incoming = _make_workflow_def([_qa_node(api_key=masked_val)])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["data"]["qa_api_key"] == real_key

    def test_new_key_is_accepted(self):
        """When user provides a brand new key, it replaces the old one."""
        old_key = "sk-proj-abcdefghijklmnop"
        new_key = "sk-proj-zyxwvutsrqponmlk"

        existing = _make_workflow_def([_qa_node(api_key=old_key)])
        incoming = _make_workflow_def([_qa_node(api_key=new_key)])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["data"]["qa_api_key"] == new_key

    def test_no_existing_qa_node(self):
        """New QA node with no prior existing node keeps incoming key."""
        new_key = "sk-brand-new-key1234"

        existing = _make_workflow_def([_agent_node()])
        incoming = _make_workflow_def([_qa_node(api_key=new_key)])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["data"]["qa_api_key"] == new_key

    def test_masked_tuner_key_is_restored(self):
        """Masked Tuner keys round-trip without losing the stored secret."""
        real_key = "tuner_live_abcdefghijklmnop"
        existing = _make_workflow_def([_tuner_node(api_key=real_key)])
        incoming = _make_workflow_def([_tuner_node(api_key=mask_key(real_key))])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["data"]["tuner_api_key"] == real_key

    def test_no_incoming_api_key(self):
        """QA node without api_key in incoming is left alone."""
        existing = _make_workflow_def([_qa_node(api_key="sk-existing-key1")])
        incoming = _make_workflow_def([_qa_node()])

        result = merge_workflow_api_keys(incoming, existing)

        assert "qa_api_key" not in result["nodes"][0]["data"]

    def test_multiple_qa_nodes_matched_by_id(self):
        """Multiple QA nodes are matched by node ID, not position."""
        key_1 = "sk-first-key-abcd1234"
        key_2 = "sk-second-key-efgh5678"

        existing = _make_workflow_def(
            [
                _qa_node(node_id="qa-1", api_key=key_1),
                _qa_node(node_id="qa-2", api_key=key_2),
            ]
        )
        incoming = _make_workflow_def(
            [
                _qa_node(node_id="qa-2", api_key=mask_key(key_2)),
                _qa_node(node_id="qa-1", api_key=mask_key(key_1)),
            ]
        )

        result = merge_workflow_api_keys(incoming, existing)

        node_map = {n["id"]: n for n in result["nodes"]}
        assert node_map["qa-1"]["data"]["qa_api_key"] == key_1
        assert node_map["qa-2"]["data"]["qa_api_key"] == key_2

    def test_none_incoming_returns_none(self):
        existing = _make_workflow_def([_qa_node(api_key="sk-key")])
        assert merge_workflow_api_keys(None, existing) is None

    def test_none_existing_returns_incoming(self):
        incoming = _make_workflow_def([_qa_node(api_key="sk-key")])
        result = merge_workflow_api_keys(incoming, None)
        assert result["nodes"][0]["data"]["qa_api_key"] == "sk-key"

    def test_non_qa_nodes_not_affected(self):
        """Agent nodes pass through without modification."""
        existing = _make_workflow_def([_agent_node()])
        incoming = _make_workflow_def([_agent_node()])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["type"] == "agentNode"

    def test_existing_node_has_no_key(self):
        """If existing QA node had no key, incoming key is kept."""
        new_key = "sk-new-key-abcd1234"

        existing = _make_workflow_def([_qa_node()])
        incoming = _make_workflow_def([_qa_node(api_key=new_key)])

        result = merge_workflow_api_keys(incoming, existing)

        assert result["nodes"][0]["data"]["qa_api_key"] == new_key


# ---------------------------------------------------------------------------
# Round-trip: mask then merge
# ---------------------------------------------------------------------------


class TestMaskAndMergeRoundTrip:
    def test_full_round_trip_preserves_key(self):
        """Simulates: save real key → GET masks it → PUT sends masked → merge restores."""
        real_key = "sk-proj-WZRTVpVvZEXF5s0H4y8N5n2BF6lRZhC79Zq"

        # 1. Real key stored in DB
        stored = _make_workflow_def(
            [
                _qa_node(api_key=real_key, qa_provider="openai", qa_model="gpt-4.1"),
            ]
        )

        # 2. GET response masks it
        fetched = mask_workflow_definition(stored)
        masked_key = fetched["nodes"][0]["data"]["qa_api_key"]
        assert masked_key != real_key
        assert masked_key.endswith(real_key[-4:])

        # 3. User saves without changing the key (sends masked value back)
        incoming = fetched  # same as what was fetched

        # 4. PUT merges — real key is restored
        merged = merge_workflow_api_keys(incoming, stored)
        assert merged["nodes"][0]["data"]["qa_api_key"] == real_key

    def test_round_trip_with_key_change(self):
        """User changes the key mid-round-trip — new key is accepted."""
        old_key = "sk-old-key-abcdefgh"
        new_key = "sk-new-key-zyxwvuts"

        stored = _make_workflow_def([_qa_node(api_key=old_key)])
        fetched = mask_workflow_definition(stored)

        # User replaces the masked key with a new one
        fetched["nodes"][0]["data"]["qa_api_key"] = new_key

        merged = merge_workflow_api_keys(fetched, stored)
        assert merged["nodes"][0]["data"]["qa_api_key"] == new_key