AgentLab/backend/tests/test_masking_tool_boundary.py

"""工具边界先遮（engine/masking.py 的 mask_query_result）的收口：认不准就整份遮，多遮也不改写没设遮罩的数值证据。

- 遮罩分析本身出了意外（认遮罩列、标签列、逐格换值的哪一步都算）：整份结果的每一格都是「已遮罩」，交给模型的、落库的
  都一样，masked_at_query 记全部的列。兜底不再经出意外的那一步（以前兜底又调了一遍逐格换值的那一步，它出意外时兜底跟着
  出意外，结果没交出去，也没有测试证明「整份遮」这句保证）
- 按值遮只换文字格：没设遮罩的数值列（访问次数正好等于某一笔被遮的金额）原样，存下的证据和数据库里的一样，双盲复算照常
  取得到；文字格里整格就是遮罩值、夹着遮罩值的照换，换过的格子记进快照（masked_cells），格子引用指到它们照实说是遮罩。
  占位的短词（未知、暂无……）不当遮罩值，没设遮罩的列、备注里的同一个词照原样

模型一律是 MockChatModel 剧本，数据源是现造的 SQLite 文件，手机号、姓名、金额、访问次数都是虚构的。
"""
from __future__ import annotations

import json

import pytest
from httpx import ASGITransport, AsyncClient

from app.core import artifact_store
from app.data.engine import engines
from app.engine import masking
from app.main import app
from tests.test_verify_agent import Script, q
from tests.test_verify_agent_masking import (  # noqa: F401 - clean_spend 是自动用上的夹具，masked_shop 是夹具
    MASKED,
    _clean,
    _source,
    ask,
    clean_spend,
    masked_shop,
)


# --------------------------------------------------------------------------
# 分析出意外：整份遮
# --------------------------------------------------------------------------


def _broken(*args, **kwargs):
    raise RuntimeError("遮罩分析里的意外")


PAYLOAD = {"sql": "SELECT id, buyer_phone, note FROM sales", "columns": ["id", "buyer_phone", "note"],
           "rows": [[1, "13900000001", "客户来电"], [2, "13900000002", "例行回访"]],
           "column_types": {"id": "number", "buyer_phone": "text", "note": "text"}}


@pytest.mark.parametrize("step", ["_derived_masks", "_snapshot_labels", "_masked_view"])
def test_a_failing_analysis_masks_every_cell(monkeypatch, step):
    """认遮罩列、认标签列、逐格换值，哪一步出了意外都整份遮：每一格都是「已遮罩」、类型记成文本，整列遮掉的是全部的列。"""
    monkeypatch.setattr(masking, step, _broken)
    done = masking.mask_query_result(json.loads(json.dumps(PAYLOAD)), ["buyer_phone"], "sqlite")
    assert done.payload["rows"] == [[MASKED] * 3, [MASKED] * 3]
    assert done.payload["column_types"] == {"id": "text", "buyer_phone": "text", "note": "text"}
    assert done.columns == ["id", "buyer_phone", "note"]
    assert done.payload["sql"] == PAYLOAD["sql"], "SQL 是模型自己写的，照原样"
    assert not done.secrets.strong, "认不准，不收值"
    assert done.whole, "快照记下是整份遮的"


def test_a_failing_analysis_masks_rows_of_any_shape(monkeypatch):
    """行是字典的结果同样整份遮。"""
    monkeypatch.setattr(masking, "_masked_view", _broken)
    payload = {"sql": "SELECT id, buyer_phone FROM sales", "columns": ["id", "buyer_phone"],
               "rows": [{"id": 1, "buyer_phone": "13900000001"}]}
    done = masking.mask_query_result(payload, ["buyer_phone"], "sqlite")
    assert done.payload["rows"] == [{"id": MASKED, "buyer_phone": MASKED}]


@pytest.mark.parametrize("step", ["_derived_masks", "_masked_view"])
async def test_the_model_and_the_snapshot_see_only_masks_when_the_analysis_fails(monkeypatch, masked_shop, step):
    """端到端：分析出了意外，交给模型的工具结果、落库的快照、通用接口给的都只有「已遮罩」，masked_at_query 是全部的列；
    格子引用照实说写法认不全、整份遮了，不采信——金额在这里没设遮罩，不说成「取自设置了遮罩的列」。"""
    monkeypatch.setattr(masking, step, _broken)
    name = masked_shop
    sql = "SELECT id, region, buyer_phone AS 手机号, amount FROM sales ORDER BY id"
    script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "amount"},
                                                        "rationale": "x"})
    out, _ = await ask(name)
    seen = json.loads(script.tool_results()["c1"].split("\n", 1)[1])
    assert seen["rows"] == [[MASKED] * 4] * 4, seen
    snapshot = artifact_store.load(out.queries[0]["artifact"])
    assert snapshot["rows"] == [[MASKED] * 4] * 4
    assert snapshot["masked_at_query"] == ["id", "region", "手机号", "amount"] and snapshot["masked_whole"] is True
    assert "masked_whole" not in seen, "记录只在快照里，不交给模型"
    assert _clean(json.dumps(snapshot, ensure_ascii=False))
    async with AsyncClient(transport=ASGITransport(app=app), base_url="http://test") as client:
        got = await client.get(f"/api/artifacts/{out.queries[0]['artifact']}")
    assert got.status_code == 200 and _clean(got.text) and "1200" not in got.text
    assert out.result["value"] is None
    assert out.cells[0]["reason"] == "Q1 的写法认不全，整份结果查询时已遮罩，不采信"


# --------------------------------------------------------------------------
# 多遮不改写没设遮罩的数值证据；占位的短词不当遮罩值
# --------------------------------------------------------------------------


#: 金额设了遮罩，都是五位数；访问次数没设遮罩，第二、三单正好等于别的订单的金额（20000、30000）。联系表的访问次数存成
#: 文字（'20000'），客户名设了遮罩、里面有占位的「未知」，区域和备注里也写着「未知」
BIG = ("CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, amount INTEGER, visits INTEGER, buyer_name TEXT,"
       " note TEXT);"
       "INSERT INTO sales VALUES (1, '华东', 10000, 25000, '张伟', '原因未知'), (2, '华东', 12000, 20000, '未知', ''),"
       " (3, '华南', 20000, 30000, '李娜', '客户未知来源'), (4, '未知', 30000, 22000, '王芳', ''),"
       " (5, '华北', 50000, 9000, '刘洋', '');"
       "CREATE TABLE contacts (id INTEGER PRIMARY KEY, visits_text TEXT, remark TEXT);"
       "INSERT INTO contacts VALUES (1, '20000', '上月消费 30000 元'), (2, '8', '例行回访');")


@pytest.fixture
async def big_shop(tmp_path):
    name, source_id = await _source(tmp_path, {"mask_columns": "amount, buyer_name"}, script=BIG)
    yield name
    await engines.invalidate(source_id)


async def test_an_unmasked_numeric_column_is_kept_even_when_it_equals_a_masked_value(monkeypatch, big_shop):
    """访问次数没设遮罩，第二、三单正好等于别的订单被遮的金额：交给模型的、落库的都是数据库里的原数，类型照旧是数，
    引用它照常取值（双盲复算取得到 20000）。以前按值整格换成「已遮罩」，快照里不留记录，引用时说成「不是数值」。"""
    name = big_shop
    sql = "SELECT id, amount, visits FROM sales ORDER BY id"
    script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 1, "column": "visits"},
                                                        "rationale": "x"})
    out, _ = await ask(name)
    seen = json.loads(script.tool_results()["c1"].split("\n", 1)[1])
    assert seen["rows"] == [[1, MASKED, 25000], [2, MASKED, 20000], [3, MASKED, 30000], [4, MASKED, 22000],
                            [5, MASKED, 9000]]
    snapshot = artifact_store.load(out.queries[0]["artifact"])
    assert snapshot["rows"] == seen["rows"] and snapshot["column_types"]["visits"] == "number"
    assert snapshot["masked_at_query"] == ["amount"] and "masked_cells" not in snapshot
    assert out.result["value"] == 20000 and out.cells[0]["status"] == "resolved"


async def test_text_cells_masked_by_value_are_recorded_and_cited_as_masked(monkeypatch, big_shop):
    """存成文字的数（'20000'）、备注里夹着的金额（30000）照样按值换：交给模型的、落库的都换掉，换过的格子记进快照
    （masked_cells），引用到它们照实说是遮罩，不说成「不是数值」。"""
    name = big_shop
    steps = [q(name, "SELECT amount FROM sales ORDER BY id", "c1"),
             q(name, "SELECT id, visits_text, remark FROM contacts ORDER BY id", "c2")]
    script = Script(monkeypatch, steps, lambda messages: {"value": {"alias": "Q2", "row": 0, "column": "visits_text"},
                                                          "rationale": "x"})
    out, _ = await ask(name)
    seen = json.loads(script.tool_results()["c2"].split("\n", 1)[1])
    assert seen["rows"] == [[1, MASKED, f"上月消费 {MASKED} 元"], [2, "8", "例行回访"]]
    assert "masked_cells" not in seen, "记录只在快照里，不交给模型"
    snapshot = artifact_store.load(out.queries[1]["artifact"])
    assert snapshot["rows"] == seen["rows"] and snapshot["masked_cells"] == [[0, 1], [0, 2]]
    assert snapshot["masked_at_query"] == []
    assert out.result["value"] is None
    assert out.cells[0]["reason"] == "Q2 第 0 行「visits_text」与设置了遮罩的值相同，查询时已遮罩，不采信"


async def test_a_partly_masked_text_cell_is_cited_as_masked(monkeypatch, big_shop):
    """引用指到文字里换掉了一截的格子：同样照实说是遮罩。"""
    name = big_shop
    steps = [q(name, "SELECT amount FROM sales ORDER BY id", "c1"),
             q(name, "SELECT id, remark FROM contacts ORDER BY id", "c2")]
    Script(monkeypatch, steps, {"value": {"alias": "Q2", "row": 0, "column": "remark"}, "rationale": "x"})
    out, _ = await ask(name)
    assert out.result["value"] is None
    assert out.cells[0]["reason"] == "Q2 第 0 行「remark」里有与设置了遮罩的值相同的部分，查询时已遮罩，不采信"


async def test_placeholder_words_are_not_masked_elsewhere(monkeypatch, big_shop):
    """客户名设了遮罩，里面有占位的「未知」：客户名整列照旧遮；没设遮罩的区域里的「未知」、备注里的「原因未知」「客户
    未知来源」照原样。以前「未知」当遮罩值收了，这几格都被换掉。"""
    name = big_shop
    sql = "SELECT id, buyer_name, region, note FROM sales ORDER BY id"
    script = Script(monkeypatch, [q(name, sql, "c1")], {"value": {"alias": "Q1", "row": 0, "column": "id"},
                                                        "rationale": "x"})
    out, _ = await ask(name)
    seen = json.loads(script.tool_results()["c1"].split("\n", 1)[1])
    assert [row[1] for row in seen["rows"]] == [MASKED] * 5
    assert [row[2:] for row in seen["rows"]] == [["华东", "原因未知"], ["华东", ""], ["华南", "客户未知来源"],
                                                 ["未知", ""], ["华北", ""]]
    assert out.masked_values.hit("张伟") and not out.masked_values.hit("未知")


def test_placeholders_are_never_collected():
    secrets = masking._Secrets()
    secrets.add({"columns": ["name"], "rows": [["未知"], ["N/A"], ["暂无"], ["张伟"], ["NULL"]]}, {"name"})
    assert secrets.strong == {"张伟"} and secrets.weak == {"张伟"}
    assert secrets.scrub("原因未知，暂无记录，N/A，张伟") == f"原因未知，暂无记录，N/A，{MASKED}"


# --------------------------------------------------------------------------
# 分档写在派生来源里、外层按它分组：标签列照常显示，标签不当遮罩值
# --------------------------------------------------------------------------


LAYOUT = ["id", "region", "amount", "visits"]
#: 探查得到的表结构：sales 里没有「档位」这一列
SCHEMA = {"tables": {"sales": {"qualified": "sales", "schema": None, "columns": [{"name": c} for c in LAYOUT]}}}
BAND = "CASE WHEN amount >= 20000 THEN '大单' ELSE '小单' END"
GROUPED = [
    pytest.param(f"SELECT 档位, COUNT(*) AS 笔数 FROM (SELECT *, {BAND} AS 档位 FROM sales) GROUP BY 档位 ORDER BY 档位",
                 id="子查询里 * 加分档"),
    pytest.param(f"SELECT 档位, COUNT(*) AS 笔数 FROM (SELECT s.*, {BAND} AS 档位 FROM sales s) t GROUP BY 档位 "
                 "ORDER BY 档位", id="子查询里 s.* 加分档"),
    pytest.param(f"WITH b AS (SELECT *, {BAND} AS 档位 FROM sales) SELECT 档位, COUNT(*) AS 笔数 FROM b GROUP BY 档位 "
                 "ORDER BY 档位", id="WITH 里 * 加分档"),
    pytest.param(f"WITH b AS (SELECT {BAND} AS 档位 FROM sales) SELECT 档位, COUNT(*) AS 笔数 FROM b GROUP BY 档位 "
                 "ORDER BY 档位", id="WITH 里只有分档"),
    pytest.param(f"WITH b AS (SELECT {BAND} AS 档位 FROM sales) SELECT b.档位, COUNT(*) AS 笔数 FROM b GROUP BY 1 "
                 "ORDER BY 1", id="按位置分组"),
    pytest.param(f"WITH b AS (SELECT {BAND} AS 档位 FROM sales), c AS (SELECT 档位, COUNT(*) AS 笔数 FROM b "
                 "GROUP BY 档位) SELECT * FROM c ORDER BY 档位", id="分组在中间一层"),
]
ROWS = [["大单", 3], ["小单", 2]]


@pytest.mark.parametrize("sql", GROUPED)
def test_a_band_grouped_outside_its_derived_source_is_shown(sql):
    """分档（按遮罩的金额写死的标签）写在子查询、WITH 里，外层按它分组：每档一行，档位列照常显示，标签不当遮罩值收；
    以前整列遮成「已遮罩」，模型分不清哪个数属于哪一档。"""
    payload = {"sql": sql, "columns": ["档位", "笔数"], "rows": [list(r) for r in ROWS],
               "column_types": {"档位": "text", "笔数": "number"}}
    done = masking.mask_query_result(payload, ["amount"], "sqlite", schema=SCHEMA)
    assert done.payload["rows"] == ROWS and done.columns == []
    assert not done.secrets.strong
    hidden, names = masking._derived_hidden({**payload, "rows": ROWS}, {"amount"}, "sqlite",
                                            masking._schema_columns(SCHEMA))
    assert "档位" not in hidden and "档位" in names, "SQL 里照旧认得出它取自金额"


@pytest.mark.parametrize("sql", [
    pytest.param(f"SELECT id, 档位 FROM (SELECT id, {BAND} AS 档位 FROM sales) ORDER BY id", id="逐行的档位"),
    pytest.param(f"SELECT 档位, COUNT(*) AS 笔数 FROM (SELECT {BAND} AS 档位 FROM sales UNION ALL "
                 "SELECT region FROM sales) GROUP BY 档位", id="另一支是区域"),
    pytest.param(f"SELECT 档位, COUNT(*) AS 笔数 FROM (SELECT id, {BAND} AS 档位 FROM sales) x "
                 "FULL JOIN (SELECT id, amount AS 档位 FROM sales) y USING (档位) GROUP BY 档位", id="按档位连起来"),
])
def test_a_band_that_is_not_one_row_per_band_is_still_masked(sql):
    """逐行的档位（看得出每一笔落在哪一档）、和别的值混在一列里、按它连起来又不写表名的（外连接时取的可能是另一边的
    金额），照旧整列遮。"""
    payload = {"sql": sql, "columns": ["档位", "笔数"] if "笔数" in sql else ["id", "档位"],
               "rows": [["大单", 3], ["小单", 2]] if "笔数" in sql else [[1, "小单"], [3, "大单"]]}
    done = masking.mask_query_result(payload, ["amount"], "sqlite", schema=SCHEMA)
    assert "档位" in done.columns


def test_a_grouped_band_with_cells_other_than_its_labels_is_still_masked():
    """格子里有标签以外的值（认错了、数的写法对不上）：照旧整列遮。"""
    sql = GROUPED[3].values[0]
    payload = {"sql": sql, "columns": ["档位", "笔数"], "rows": [["大单", 3], ["13900000001", 2]]}
    assert masking.mask_query_result(payload, ["amount"], "sqlite", schema=SCHEMA).columns == ["档位"]


def test_without_a_schema_a_star_next_to_the_band_is_still_masked():
    """不知道表结构时，* 展开的表可能就有「档位」这一列（外层取到的是它）：照旧整列遮，标签照收（多遮）。"""
    sql = GROUPED[0].values[0]
    payload = {"sql": sql, "columns": ["档位", "笔数"], "rows": [list(r) for r in ROWS]}
    done = masking.mask_query_result(payload, ["amount"], "sqlite")
    assert done.columns == ["档位"] and done.payload["rows"] == [[MASKED, 3], [MASKED, 2]]


# --------------------------------------------------------------------------
# 数值格按位置认：同名的列、写成字符串的数
# --------------------------------------------------------------------------


def _known(*values: str) -> masking._Secrets:
    """此前认出的遮罩值（前面的查询遮掉的、报告快照里的）。"""
    secrets = masking._Secrets()
    secrets.add({"columns": ["v"], "rows": [[v] for v in values]}, {"v"})
    return secrets


def test_a_text_column_sharing_its_name_with_a_number_column_is_still_masked_by_value():
    """column_types 按列名记，同名的几列只留下后面那一列的类型：两张表里都叫 phone，前一列是文字（写着号码）、后一列是
    整数，记成 number。以前按列名认数值列，文字列也当成了数值列，号码原样交给模型、落库；现在同名的几列按格子本身认，
    文字格照换、记进 masked_cells，真正的整数照原样。写成字符串的数在同名列里也按文字处理（宁可多遮）。"""
    payload = {"sql": "SELECT a.region, a.phone, b.phone FROM contacts a JOIN legacy b ON a.id = b.id",
               "columns": ["region", "phone", "phone"],
               "rows": [["华东", "联系人 13900000001", 13900000002], ["华南", "13900000001", "13900000001"]],
               "column_types": {"region": "text", "phone": "number"}}
    done = masking.mask_query_result(payload, [], "sqlite", _known("13900000001"))
    assert done.payload["rows"] == [["华东", f"联系人 {MASKED}", 13900000002], ["华南", MASKED, MASKED]]
    assert done.cells == [[0, 1], [1, 1], [1, 2]]


def test_a_numeric_column_written_as_decimal_strings_is_kept():
    """Postgres 的 NUMERIC、MySQL 的 DECIMAL 交回来是字符串（"20000.00"），列类型记着 number：没设遮罩的这一列正好等于某个
    遮罩值（金额也是 NUMERIC）也照原样，不记进 masked_cells——数值证据不改写。"""
    payload = {"sql": "SELECT id, visits FROM sales2 ORDER BY id", "columns": ["id", "visits"],
               "rows": [[1, "20000.00"], [2, "9000.00"]], "column_types": {"id": "number", "visits": "number"}}
    done = masking.mask_query_result(payload, [], "postgresql", _known("20000.00", "30000.00"))
    assert done.payload["rows"] == [[1, "20000.00"], [2, "9000.00"]] and done.cells == []


#: 两张表里都有 remark：工单的是文字（写着号码），评分的是整数
DUPLICATES = ("CREATE TABLE sales (id INTEGER PRIMARY KEY, region TEXT, buyer_phone TEXT, amount INTEGER);"
              "INSERT INTO sales VALUES (1, '华东', '13900000001', 1200), (2, '华南', '13900000002', 900);"
              "CREATE TABLE tickets (id INTEGER PRIMARY KEY, remark TEXT);"
              "INSERT INTO tickets VALUES (1, '客户 13900000001 来电投诉'), (2, '回访 13900000002');"
              "CREATE TABLE scores (id INTEGER PRIMARY KEY, remark INTEGER);"
              "INSERT INTO scores VALUES (1, 5), (2, 4);")


@pytest.fixture
async def ticket_shop(tmp_path):
    name, source_id = await _source(tmp_path, {"mask_columns": "buyer_phone"}, script=DUPLICATES)
    yield name
    await engines.invalidate(source_id)


@pytest.mark.parametrize("sql", [
    "SELECT t.remark, s.remark FROM tickets t JOIN scores s ON s.id = t.id ORDER BY t.id",
    "SELECT * FROM tickets t JOIN scores s ON s.id = t.id ORDER BY t.id",
])
async def test_duplicate_names_do_not_let_numbers_shield_text(monkeypatch, ticket_shop, sql):
    """端到端：先查出号码（查询时遮掉），再把工单和评分连起来查，两列都叫 remark（* 连两张表时 id 也同名）。工单备注里的
    号码交给模型、落库之前照样换掉，评分照原样；以前评分的 number 套到了同名的备注上，号码原样交了出去。"""
    name = ticket_shop
    steps = [q(name, "SELECT id, buyer_phone FROM sales ORDER BY id", "c1"), q(name, sql, "c2")]
    script = Script(monkeypatch, steps, {"value": None, "rationale": "x"})
    out, _ = await ask(name)
    seen = script.tool_results()["c2"]
    assert "1390000000" not in seen and f"客户 {MASKED} 来电投诉" in seen, seen
    snapshot = artifact_store.load(out.queries[1]["artifact"])
    remarks = [i for i, c in enumerate(snapshot["columns"]) if c == "remark"]
    assert [[row[i] for i in remarks] for row in snapshot["rows"]] == [[f"客户 {MASKED} 来电投诉", 5],
                                                                      [f"回访 {MASKED}", 4]]
    assert snapshot["masked_cells"] == [[0, remarks[0]], [1, remarks[0]]]


async def test_a_decimal_column_is_cited_as_a_number(monkeypatch, composite_numbers):
    """端到端：NUMERIC 的访问次数（字符串 "20000.00"、列类型 number）正好等于某一笔被遮的金额：交给模型的、落库的照原样，
    格子引用照常取到 20000（双盲复算拿得到它）。"""
    name, sql = composite_numbers
    script = Script(monkeypatch, [q(name, "SELECT id, amount FROM sales ORDER BY id", "c1"), q(name, sql, "c2")],
                    {"value": {"alias": "Q2", "row": 0, "column": "visits"}, "rationale": "x"})
    out, _ = await ask(name)
    seen = json.loads(script.tool_results()["c2"].split("\n", 1)[1])
    assert seen["rows"] == [[1, "20000.00"], [2, "9000.00"]]
    snapshot = artifact_store.load(out.queries[1]["artifact"])
    assert snapshot["rows"] == seen["rows"] and "masked_cells" not in snapshot
    assert out.result["value"] == 20000 and out.cells[0]["status"] == "resolved"


@pytest.fixture
async def composite_numbers(monkeypatch):
    """造一个 Postgres 数据源（查询交给换掉的 run_query，不连库）：金额设了遮罩，金额、访问次数都是 NUMERIC，驱动交回
    Decimal，查询层转成字符串、列类型记 number。交回 (源名, 查访问次数的 SQL)。"""
    import uuid

    from app.data.engine import QueryResult
    from app.db.base import SessionLocal
    from app.db.models import DataSource
    from app.tools import datasource

    visits = "SELECT id, visits FROM sales ORDER BY id"
    results = {"SELECT id, amount FROM sales ORDER BY id": (["id", "amount"], [[1, "20000.00"], [2, "30000.00"]]),
               visits: (["id", "visits"], [[1, "20000.00"], [2, "9000.00"]])}

    async def driver(source, statement, **kwargs):
        columns, rows = results[statement]
        return QueryResult(columns=columns, rows=[list(r) for r in rows], row_count=len(rows), truncated=False,
                           elapsed_ms=1, sql=statement, column_types={c: "number" for c in columns})

    monkeypatch.setattr(datasource, "run_query", driver)
    name = f"pg_{uuid.uuid4().hex[:8]}"
    async with SessionLocal() as session:
        row = DataSource(name=name, kind="postgresql", host="127.0.0.1", port=5432, username="probe", database="shop",
                         readonly=True, enabled=True, options={"mask_columns": "amount"},
                         schema_cache={"tables": {"sales": {"columns": [{"name": c} for c in ("id", "amount",
                                                                                             "visits")]}}})
        session.add(row)
        await session.commit()
        source_id = row.id
    yield name, visits
    await engines.invalidate(source_id)