tests/test_thoughtspot_deployer_errors.py

from thoughtspot_errors import friendly_mcp_liveboard_error
from tests.e2e_quality import extract_db_schema_from_tml_docs
import liveboard_creator
import thoughtspot_deployer


def test_friendly_mcp_error_classifies_service_outage():
    message, category = friendly_mcp_liveboard_error(
        "MCP liveboard timed out after 360s — MCP endpoint unresponsive at network level"
    )

    assert category == "mcp_service"
    assert "MCP service appears unavailable or timed out" in message
    assert "Raw MCP error:" in message


def test_friendly_mcp_error_classifies_auth_failure():
    message, category = friendly_mcp_liveboard_error(
        "Failed to get trusted auth token. Check THOUGHTSPOT_TRUSTED_AUTH_KEY"
    )

    assert category == "mcp_auth"
    assert "could not authenticate" in message
    assert "trusted auth key" in message.lower()


def test_friendly_mcp_error_classifies_generation_failure():
    message, category = friendly_mcp_liveboard_error(
        "No questions generated. Model may be empty or incompatible."
    )

    assert category == "mcp_generation"
    assert "MCP liveboard generation step" in message
    assert "falling back" in message


def test_friendly_mcp_error_classifies_missing_answer_tokens_as_recoverable():
    message, category = friendly_mcp_liveboard_error(
        "No liveboard-compatible answers with ThoughtSpot session/tokens were retrieved."
    )

    assert category == "mcp_answer_tokens"
    assert "backup liveboard" in message


def test_direct_answer_rejects_empty_token_payload(monkeypatch):
    class FakeResponse:
        status_code = 200
        text = ""

        def json(self):
            return {
                "session_identifier": "session-123",
                "generation_number": -1,
                "message_type": "TSAnswer",
                "visualization_type": "Undefined",
                "tokens": "",
                "display_tokens": "",
            }

    class FakeSession:
        def post(self, *args, **kwargs):
            return FakeResponse()

    monkeypatch.setattr(liveboard_creator, "_get_direct_api_session", lambda *args, **kwargs: (FakeSession(), "token"))

    assert liveboard_creator._get_answer_direct(
        "Total Net Sales",
        "model-guid",
        "testrunner@example.com",
        "secret",
        ts_url="https://example.thoughtspot.cloud",
    ) is None


def test_liveboard_answer_validation_requires_tokens_but_allows_undefined_visualization():
    assert liveboard_creator._is_liveboard_answer_valid(
        {
            "session_identifier": "session-123",
            "tokens": [{"token": "Net Sales"}],
            "visualization_type": "Undefined",
        }
    )
    assert not liveboard_creator._is_liveboard_answer_valid(
        {
            "session_identifier": "session-123",
            "tokens": "",
            "visualization_type": "Undefined",
        }
    )


def test_extract_db_schema_reads_physical_table_tml_not_model_name():
    docs = [
        """
model:
  name: TST_WRONG_01010000_ABC_mdl
""",
        """
table:
  name: SALES
  db: DEMOBUILD
  schema: TST_RIGHT_01010000_DEF_sch
""",
    ]

    assert extract_db_schema_from_tml_docs(docs) == ("DEMOBUILD", "TST_RIGHT_01010000_DEF_sch")


def test_extract_db_schema_does_not_guess_from_model_name():
    docs = [
        """
model:
  name: TST_WRONG_01010000_ABC_mdl
"""
    ]

    assert extract_db_schema_from_tml_docs(docs) == ("", "")


def test_connection_create_reconciles_secure_store_bad_gateway(monkeypatch):
    class FakeResponse:
        status_code = 200
        text = ""

        def json(self):
            return [
                {
                    "response": {
                        "status": {
                            "status_code": "ERROR",
                            "error_message": (
                                "Exception occurred while saving dataSource: "
                                "Error Code: SECURE_STORE_ERROR Error Message: Bad Gateway."
                            ),
                        },
                        "header": {"name": "MBB_VIZ_TEST_conn"},
                    }
                }
            ]

    class FakeSession:
        def __init__(self):
            self.posts = []

        def post(self, url, **kwargs):
            self.posts.append({"url": url, **kwargs})
            return FakeResponse()

    deployer = thoughtspot_deployer.ThoughtSpotDeployer.__new__(
        thoughtspot_deployer.ThoughtSpotDeployer
    )
    deployer.base_url = "https://ts.example"
    deployer.session = FakeSession()
    deployer.headers = {}
    deployer.create_connection_tml = lambda name, database: "connection:\n  name: MBB_VIZ_TEST_conn\n"

    calls = []

    def fake_get_connection_by_name(name):
        calls.append(name)
        if len(calls) == 1:
            return None
        return {"header": {"id_guid": "connection-guid", "name": name}}

    deployer.get_connection_by_name = fake_get_connection_by_name
    monkeypatch.setattr(thoughtspot_deployer.time, "sleep", lambda seconds: None)

    guid, fqn = deployer.create_connection_with_reconcile(
        "MBB_VIZ_TEST_conn",
        "DEMOBUILD_TEST",
        log_progress=lambda message: None,
    )

    assert guid == "connection-guid"
    assert fqn == "connection-guid"
    assert len(deployer.session.posts) == 1
    assert len(calls) == 2


def test_trusted_auth_retries_transient_gateway_error(monkeypatch):
    class FakeResponse:
        def __init__(self, status_code, payload=None, text=""):
            self.status_code = status_code
            self._payload = payload or {}
            self.text = text

        def json(self):
            return self._payload

    class FakeSession:
        def __init__(self):
            self.headers = {}
            self.calls = 0

        def post(self, url, **kwargs):
            self.calls += 1
            if self.calls == 1:
                return FakeResponse(502, text="Bad Gateway")
            return FakeResponse(200, {"token": "token-123"})

    deployer = thoughtspot_deployer.ThoughtSpotDeployer.__new__(
        thoughtspot_deployer.ThoughtSpotDeployer
    )
    deployer.base_url = "https://ts.example"
    deployer.username = "mike.boone@thoughtspot.com"
    deployer.secret_key = "secret-key-value"
    deployer.session = FakeSession()
    deployer.last_auth_status_code = None
    deployer.last_auth_error = ""

    monkeypatch.setenv("TS_AUTH_MAX_ATTEMPTS", "2")
    monkeypatch.setenv("TS_AUTH_RETRY_WAIT_SECONDS", "1")
    monkeypatch.setattr(thoughtspot_deployer.time, "sleep", lambda seconds: None)

    assert deployer.authenticate_trusted() is True
    assert deployer.session.calls == 2
    assert deployer.session.headers["Authorization"] == "Bearer token-123"