agent-management/backend/rag-ai-bridge/test_server.py

import pytest

from server import _extract_code, _read_only


@pytest.mark.parametrize(
    ("language", "raw", "expected"),
    [
        ("SQL", "<think>hidden</think>```sql\nSELECT * FROM rescue_forces\n```", "SELECT * FROM rescue_forces"),
        ("SQL", "Here is the query:\nSELECT name FROM rescue_forces\nExplanation: done", "SELECT name FROM rescue_forces"),
        ("Cypher", "```cypher\nMATCH (m:Mission) RETURN m\n```", "MATCH (m:Mission) RETURN m"),
        ("Cypher", "思考\nMATCH (s:Stage) RETURN s\n说明：done", "MATCH (s:Stage) RETURN s"),
    ],
)
def test_extract_and_accept_read_only(language, raw, expected):
    assert _read_only(_extract_code(raw, language), language) == expected


@pytest.mark.parametrize(
    ("language", "query"),
    [
        ("SQL", "DELETE FROM rescue_forces"),
        ("SQL", "SELECT 1; DROP TABLE rescue_forces"),
        ("Cypher", "MATCH (n) DELETE n"),
        ("Cypher", "MATCH (n) RETURN n; CREATE ()"),
    ],
)
def test_rejects_writes_and_multiple_statements(language, query):
    with pytest.raises(ValueError):
        _read_only(query, language)