|
| 1 | +"""Offline task regressions with model, persistence and database I/O replaced.""" |
| 2 | + |
| 3 | +from types import SimpleNamespace |
| 4 | +from unittest.mock import Mock |
| 5 | + |
| 6 | +import pytest |
| 7 | +import sqlbot_xpack # noqa: F401 # Initialize the published wheel before app imports. |
| 8 | +from pydantic import ValidationError |
| 9 | + |
| 10 | +from apps.chat.models.chat_model import ChatFinishStep |
| 11 | +from apps.chat.task import llm |
| 12 | +from common.core.config import Settings |
| 13 | + |
| 14 | + |
| 15 | +@pytest.mark.parametrize( |
| 16 | + "value,expected", |
| 17 | + [(None, False), ("false", False), ("true", True), (" TRUE ", True)], |
| 18 | +) |
| 19 | +def test_startup_configuration(monkeypatch, value, expected): |
| 20 | + monkeypatch.delenv("SQLBOT_SKIP_TABLE_CHECK", raising=False) |
| 21 | + if value is not None: |
| 22 | + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", value) |
| 23 | + config = Settings(_env_file=None) |
| 24 | + assert config.SQLBOT_SKIP_TABLE_CHECK is expected |
| 25 | + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", str(not expected)) |
| 26 | + assert config.SQLBOT_SKIP_TABLE_CHECK is expected |
| 27 | + |
| 28 | + |
| 29 | +def test_invalid_startup_configuration_is_rejected(monkeypatch): |
| 30 | + monkeypatch.setenv("SQLBOT_SKIP_TABLE_CHECK", "invalid") |
| 31 | + with pytest.raises(ValidationError, match="SQLBOT_SKIP_TABLE_CHECK"): |
| 32 | + Settings(_env_file=None) |
| 33 | + |
| 34 | + |
| 35 | +@pytest.fixture |
| 36 | +def task(monkeypatch): |
| 37 | + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", False) |
| 38 | + monkeypatch.setattr(llm, "session_maker", Mock()) |
| 39 | + monkeypatch.setattr(llm, "check_connection", Mock(return_value=True)) |
| 40 | + monkeypatch.setattr(llm, "start_log", Mock()) |
| 41 | + monkeypatch.setattr(llm, "end_log", Mock()) |
| 42 | + monkeypatch.setattr(llm, "get_chat_chart_data", Mock(return_value=[])) |
| 43 | + service = Mock(spec=llm.LLMService) |
| 44 | + service.ds = SimpleNamespace(id=1, type="pg") |
| 45 | + service.chat_question = SimpleNamespace(question="sales") |
| 46 | + service.record = SimpleNamespace(id=1) |
| 47 | + service.get_record.return_value = service.record |
| 48 | + service.current_user = SimpleNamespace(id=2, isAdmin=False, weight=0) |
| 49 | + service.current_assistant = None |
| 50 | + service.change_title = False |
| 51 | + service.table_name_list = ["sales"] |
| 52 | + service.current_logs = {} |
| 53 | + service.generate_sql.return_value = [{"content": "{}"}] |
| 54 | + service.generate_filter.return_value = None |
| 55 | + service.execute_sql.return_value = {"data": []} |
| 56 | + return service |
| 57 | + |
| 58 | + |
| 59 | +def run(task, sql, reported_tables=None): |
| 60 | + # Model-reported tables are deliberately independent of the real SQL. |
| 61 | + task.check_sql.return_value = sql, reported_tables or ["sales"] |
| 62 | + task.check_save_sql.return_value = sql |
| 63 | + results = list( |
| 64 | + llm.LLMService.run_task( |
| 65 | + task, in_chat=False, stream=False, finish_step=ChatFinishStep.QUERY_DATA |
| 66 | + ) |
| 67 | + ) |
| 68 | + assert len(results) == 1 |
| 69 | + return results[0] |
| 70 | + |
| 71 | + |
| 72 | +@pytest.mark.parametrize( |
| 73 | + "sql", |
| 74 | + [ |
| 75 | + "SELECT id FROM sales", |
| 76 | + "WITH totals AS (SELECT id FROM sales) SELECT * FROM totals", |
| 77 | + ], |
| 78 | +) |
| 79 | +def test_default_allows_known_tables_and_cte_aliases(task, sql): |
| 80 | + result = run(task, sql) |
| 81 | + assert result["success"] is True |
| 82 | + task.execute_sql.assert_called_once_with(sql=sql) |
| 83 | + task.check_save_sql.assert_called_once() |
| 84 | + task.save_error.assert_not_called() |
| 85 | + |
| 86 | + |
| 87 | +@pytest.mark.parametrize( |
| 88 | + "sql,error", |
| 89 | + [ |
| 90 | + ("SELECT * FROM secret_payroll", "unauthorized tables: secret_payroll"), |
| 91 | + ( |
| 92 | + "SELECT id FROM sales UNION ALL SELECT id FROM secret_payroll", |
| 93 | + "unauthorized tables: secret_payroll", |
| 94 | + ), |
| 95 | + ("SELECT 1", "unable to extract table names"), |
| 96 | + ("SELECT (", "unable to extract table names"), |
| 97 | + ], |
| 98 | +) |
| 99 | +def test_default_rejects_invalid_or_unknown_tables_before_save_and_execute( |
| 100 | + task, sql, error |
| 101 | +): |
| 102 | + result = run(task, sql, reported_tables=["sales"]) |
| 103 | + assert result["success"] is False |
| 104 | + assert error in result["message"] |
| 105 | + task.save_error.assert_called_once() |
| 106 | + task.check_save_sql.assert_not_called() |
| 107 | + task.execute_sql.assert_not_called() |
| 108 | + |
| 109 | + |
| 110 | +@pytest.mark.parametrize("sql", ["SELECT * FROM secret_payroll", "SELECT 1"]) |
| 111 | +def test_explicit_skip_bypasses_extraction_and_keeps_remaining_flow( |
| 112 | + task, monkeypatch, sql |
| 113 | +): |
| 114 | + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", True) |
| 115 | + extract = Mock(side_effect=AssertionError("table extraction must be skipped")) |
| 116 | + monkeypatch.setattr(llm, "extract_tables_from_sql", extract) |
| 117 | + result = run(task, sql) |
| 118 | + assert result["success"] is True |
| 119 | + extract.assert_not_called() |
| 120 | + task.generate_filter.assert_called_once() |
| 121 | + task.check_save_sql.assert_called_once() |
| 122 | + task.execute_sql.assert_called_once_with(sql=sql) |
| 123 | + task.save_error.assert_not_called() |
| 124 | + |
| 125 | + |
| 126 | +@pytest.mark.parametrize( |
| 127 | + "sql", |
| 128 | + [ |
| 129 | + "DELETE FROM sales", |
| 130 | + "SELECT pg_read_file('/etc/passwd')", |
| 131 | + "SHOW TABLES", |
| 132 | + ], |
| 133 | +) |
| 134 | +def test_skip_keeps_database_execution_guards(task, monkeypatch, sql): |
| 135 | + from apps.db import db |
| 136 | + |
| 137 | + monkeypatch.setattr(llm.settings, "SQLBOT_SKIP_TABLE_CHECK", True) |
| 138 | + monkeypatch.setattr(llm.settings, "SQLBOT_ALLOW_METADATA_QUERIES", False) |
| 139 | + connect = Mock(side_effect=AssertionError("database must not be contacted")) |
| 140 | + monkeypatch.setattr(db, "get_session", connect) |
| 141 | + task.execute_sql.side_effect = lambda sql: llm.LLMService.execute_sql(task, sql) |
| 142 | + result = run(task, sql) |
| 143 | + assert result["success"] is False |
| 144 | + assert "SQL can only contain read operations" in result["message"] |
| 145 | + task.execute_sql.assert_called_once_with(sql=sql) |
| 146 | + connect.assert_not_called() |
0 commit comments