tests/test_search_query_validation.py

"""Regression tests for liveboard search-query validation.

Covers the three query bugs from run TRI_08250424_PY3 (Triumph Financial,
liveboard ca684d71-dd94-4d70-a24a-efb2e443f9ea):
  1. Cross-fact date on a granularity token (blank KPI tile)
  2. Cross-fact date on a time filter (blank/misleading chart)
  3. Granularity finer than the fact's data grain (.weekly on monthly snapshot)

No network calls. The fixtures mirror what ThoughtSpot ACTUALLY exports:
the model TML has no 'calendar' property on date columns (TS strips it on
export), so date detection must come from the associated table TMLs'
db_column_properties.data_type — exactly what build_column_map_from_tml_objects
consumes.
"""

import copy
import json

import pytest

from liveboard_creator import (
    build_column_map_from_tml_objects,
    build_search_column_map,
    validate_search_query,
    validate_liveboard_queries,
    _apply_query_fix_to_viz,
    _infer_date_grain,
)


def _model_col(name, table, physical, ctype):
    # Real exported model TML: no 'calendar', just column_type/index_type
    return {
        'name': name,
        'column_id': f'{table}::{physical}',
        'properties': {'column_type': ctype, 'index_type': 'DONT_INDEX'},
    }


def _table_tml(table_name, col_types):
    return {'table': {
        'name': table_name,
        'db': 'DB', 'schema': 'SCHEMA', 'db_table': table_name,
        'columns': [
            {'name': col, 'db_column_name': col,
             'properties': {'column_type': 'ATTRIBUTE'},
             'db_column_properties': {'data_type': dtype}}
            for col, dtype in col_types.items()
        ],
    }}


# Mirrors the Triumph model: load transactions are daily (Load Date),
# carrier capacity is a monthly snapshot (Capacity Month).
MODEL_TML = {'guid': 'model-guid', 'model': {'name': 'TRI_mdl', 'columns': [
    _model_col('Load Date', 'FACT_LOAD_TRANSACTION', 'LOAD_DATE', 'ATTRIBUTE'),
    _model_col('Loads Tendered', 'FACT_LOAD_TRANSACTION', 'LOADS_TENDERED', 'MEASURE'),
    _model_col('Potential Savings Usd', 'FACT_LOAD_TRANSACTION', 'POTENTIAL_SAVINGS_USD', 'MEASURE'),
    _model_col('Spot Rate Per Mile', 'FACT_LOAD_TRANSACTION', 'SPOT_RATE_PER_MILE', 'MEASURE'),
    _model_col('Network Rate Per Mile', 'FACT_LOAD_TRANSACTION', 'NETWORK_RATE_PER_MILE', 'MEASURE'),
    _model_col('Capacity Month', 'FACT_CARRIER_CAPACITY', 'CAPACITY_MONTH', 'ATTRIBUTE'),
    _model_col('Capacity Utilization Rate', 'FACT_CARRIER_CAPACITY', 'CAPACITY_UTILIZATION_RATE', 'MEASURE'),
    # Attribute with a grain word in its name but an INT type — NOT a date
    _model_col('Member Since Year', 'DIM_BROKER', 'MEMBER_SINCE_YEAR', 'ATTRIBUTE'),
    _model_col('Lane Tier', 'DIM_LANE', 'LANE_TIER', 'ATTRIBUTE'),
    _model_col('Dim Broker Name', 'DIM_BROKER', 'DIM_BROKER_NAME', 'ATTRIBUTE'),
]}}

TABLE_TMLS = [
    _table_tml('FACT_LOAD_TRANSACTION', {
        'LOAD_DATE': 'DATE', 'LOADS_TENDERED': 'INT64',
        'POTENTIAL_SAVINGS_USD': 'DOUBLE', 'SPOT_RATE_PER_MILE': 'DOUBLE',
        'NETWORK_RATE_PER_MILE': 'DOUBLE',
    }),
    _table_tml('FACT_CARRIER_CAPACITY', {
        'CAPACITY_MONTH': 'DATE', 'CAPACITY_UTILIZATION_RATE': 'DOUBLE',
    }),
    _table_tml('DIM_BROKER', {'MEMBER_SINCE_YEAR': 'INT64', 'DIM_BROKER_NAME': 'VARCHAR'}),
    _table_tml('DIM_LANE', {'LANE_TIER': 'VARCHAR'}),
]


@pytest.fixture()
def column_map():
    cmap = build_column_map_from_tml_objects([MODEL_TML] + TABLE_TMLS)
    assert cmap is not None
    return cmap


def test_date_detection_comes_from_table_data_types(column_map):
    # Model TML has no 'calendar' (real export) — dates must still be found
    load_date = column_map['columns']['loaddate']
    assert load_date['is_date'] and load_date['grain'] == 'daily'
    cap_month = column_map['columns']['capacitymonth']
    assert cap_month['is_date'] and cap_month['grain'] == 'monthly'
    # INT column with a grain word in its name is NOT a date
    assert column_map['columns']['membersinceyear']['is_date'] is False
    # Fact tables carry their own date columns
    assert column_map['tables']['FACT_LOAD_TRANSACTION']['date_cols'][0]['name'] == 'Load Date'


def test_calendar_property_is_only_a_fallback():
    # Without table data types, a column is only a date if 'calendar' is present
    cols = copy.deepcopy(MODEL_TML['model']['columns'])
    assert build_search_column_map(cols)['columns']['loaddate']['is_date'] is False
    cols[0]['properties']['calendar'] = 'calendar'
    assert build_search_column_map(cols)['columns']['loaddate']['is_date'] is True


def test_grain_inference_uses_last_word_only():
    assert _infer_date_grain('Capacity Month') == 'monthly'
    assert _infer_date_grain('Month End Date') == 'daily'   # not monthly
    assert _infer_date_grain('Fiscal Year') == 'yearly'
    assert _infer_date_grain('Snapshot Week') == 'weekly'
    assert _infer_date_grain('Load Date') == 'daily'


def test_cross_fact_date_on_granularity_token(column_map):
    # Triumph bug 1: measure on load fact, date from capacity fact -> blank tile
    fixed, corrections = validate_search_query(
        "sum [Potential Savings Usd] [Capacity Month].monthly", column_map)
    assert fixed == "sum [Potential Savings Usd] [Load Date].monthly"
    assert len(corrections) == 1
    assert corrections[0]['old_col'] == 'Capacity Month'
    assert corrections[0]['new_col'] == 'Load Date'


def test_cross_fact_date_on_filter_token(column_map):
    # Triumph bug 2: load-fact measures filtered by the capacity fact's date
    fixed, corrections = validate_search_query(
        "[Spot Rate Per Mile] [Network Rate Per Mile] [Lane Tier] "
        "[Capacity Month] = 'last 18 months'", column_map)
    assert fixed == ("[Spot Rate Per Mile] [Network Rate Per Mile] [Lane Tier] "
                     "[Load Date] = 'last 18 months'")
    assert corrections and corrections[0]['old_gran'] is None


def test_granularity_finer_than_data_grain(column_map):
    # Triumph bug 3: .weekly on a monthly snapshot fact
    fixed, corrections = validate_search_query(
        "average [Capacity Utilization Rate] [Capacity Month].weekly", column_map)
    assert fixed == "average [Capacity Utilization Rate] [Capacity Month].monthly"
    assert corrections[0]['old_gran'] == 'weekly'
    assert corrections[0]['new_gran'] == 'monthly'


def test_valid_queries_pass_untouched(column_map):
    for query in [
        "sum [Loads Tendered] [Load Date].monthly",
        "average [Capacity Utilization Rate] [Capacity Month].monthly",
        "sum [Loads Tendered] [Load Date].weekly [Load Date].'last 8 quarters'",
        "top 10 [Dim Broker Name] [Load Date] = 'last 12 months' sort by [Potential Savings Usd]",
    ]:
        fixed, corrections = validate_search_query(query, column_map)
        assert fixed == query
        assert corrections == []


def test_mixed_fact_measures_are_left_alone(column_map):
    # Measures from BOTH facts: no single "right" date column exists, so the
    # validator must not rewrite anything (rewriting would just move the
    # cross-fact problem onto the other measure).
    query = ("sum [Potential Savings Usd] [Capacity Utilization Rate] "
             "[Capacity Month].monthly")
    fixed, corrections = validate_search_query(query, column_map)
    assert fixed == query
    assert corrections == []


def test_kpi_double_date_usage_replaces_both_tokens(column_map):
    # KPI pattern: same date appears twice (granularity + time filter)
    fixed, corrections = validate_search_query(
        "sum [Potential Savings Usd] [Capacity Month].monthly "
        "[Capacity Month].'last 8 quarters'", column_map)
    assert fixed == ("sum [Potential Savings Usd] [Load Date].monthly "
                     "[Load Date].'last 8 quarters'")


def test_query_without_known_measures_is_left_alone(column_map):
    fixed, corrections = validate_search_query(
        "[Dim Broker Name] [Capacity Month].weekly", column_map)
    assert fixed == "[Dim Broker Name] [Capacity Month].weekly"
    assert corrections == []


def _kpi_viz(search_query, date_label, measure_label, name):
    """A minimal exported-TML KPI viz with the dependent column labels."""
    return {
        'id': 'Viz_1',
        'answer': {
            'name': name,
            'tables': [{'id': 'mdl', 'name': 'mdl'}],
            'search_query': search_query,
            'answer_columns': [{'name': date_label}, {'name': measure_label}],
            'table': {
                'table_columns': [
                    {'column_id': date_label, 'headline_aggregation': 'MIN-MAX'},
                    {'column_id': measure_label, 'headline_aggregation': 'SUM'},
                ],
                'ordered_column_ids': [date_label, measure_label],
            },
            'chart': {
                'type': 'KPI',
                'chart_columns': [{'column_id': date_label}, {'column_id': measure_label}],
                'axis_configs': [{'x': [date_label], 'y': [measure_label]}],
                'client_state_v2': json.dumps({
                    'axisProperties': [{'properties': {'linkedColumns': [date_label]}}]
                }),
            },
        },
    }


def test_apply_fix_rewrites_dependent_labels(column_map):
    viz = _kpi_viz(
        "sum [Potential Savings Usd] [Capacity Month].monthly",
        'Month(Capacity Month)', 'Total Potential Savings Usd',
        'Total Potential Savings Monthly',
    )
    tml = {'liveboard': {'visualizations': [copy.deepcopy(viz)]}}
    fixed_count, messages = validate_liveboard_queries(tml, column_map)
    assert fixed_count == 1

    answer = tml['liveboard']['visualizations'][0]['answer']
    assert answer['search_query'] == "sum [Potential Savings Usd] [Load Date].monthly"
    assert answer['answer_columns'][0]['name'] == 'Month(Load Date)'
    assert answer['table']['ordered_column_ids'][0] == 'Month(Load Date)'
    assert answer['chart']['chart_columns'][0]['column_id'] == 'Month(Load Date)'
    assert answer['chart']['axis_configs'][0]['x'] == ['Month(Load Date)']
    assert 'Month(Load Date)' in answer['chart']['client_state_v2']
    assert 'Capacity Month' not in json.dumps(answer)


def test_apply_fix_updates_title_on_granularity_change(column_map):
    viz = _kpi_viz(
        "average [Capacity Utilization Rate] [Capacity Month].weekly",
        'Week(Capacity Month)', 'Average Capacity Utilization Rate',
        'Average Capacity Utilization Rate Weekly',
    )
    tml = {'liveboard': {'visualizations': [viz]}}
    fixed_count, _ = validate_liveboard_queries(tml, column_map)
    assert fixed_count == 1

    answer = tml['liveboard']['visualizations'][0]['answer']
    assert answer['search_query'] == "average [Capacity Utilization Rate] [Capacity Month].monthly"
    assert answer['name'] == 'Average Capacity Utilization Rate Monthly'
    assert answer['answer_columns'][0]['name'] == 'Month(Capacity Month)'


def test_apply_fix_bails_out_when_labels_missing(column_map):
    # Dependent label absent from the answer -> viz must be left untouched
    viz = _kpi_viz(
        "sum [Potential Savings Usd] [Capacity Month].monthly",
        'SomethingUnexpected', 'Total Potential Savings Usd',
        'Total Potential Savings Monthly',
    )
    original = copy.deepcopy(viz)
    fixed_query, corrections = validate_search_query(
        viz['answer']['search_query'], column_map)
    assert corrections
    assert _apply_query_fix_to_viz(viz, fixed_query, corrections) is False
    assert viz == original


def test_error_on_one_viz_does_not_abandon_the_rest(column_map):
    # First viz has a non-JSON-serializable value; second is a normal fixable
    # KPI. The bad viz is reported and skipped, the good one still gets fixed.
    bad = _kpi_viz(
        "sum [Potential Savings Usd] [Capacity Month].monthly",
        'Month(Capacity Month)', 'Total Potential Savings Usd', 'Bad Viz',
    )
    bad['answer']['chart']['weird'] = object()  # json.dumps will raise
    good = _kpi_viz(
        "sum [Potential Savings Usd] [Capacity Month].monthly",
        'Month(Capacity Month)', 'Total Potential Savings Usd', 'Good Viz',
    )
    good['id'] = 'Viz_2'
    tml = {'liveboard': {'visualizations': [bad, good]}}
    fixed_count, messages = validate_liveboard_queries(tml, column_map)
    assert fixed_count == 1
    assert any('Bad Viz' in m and 'error' in m for m in messages)
    assert good['answer']['search_query'] == "sum [Potential Savings Usd] [Load Date].monthly"


def test_note_tiles_and_queryless_vizzes_are_skipped(column_map):
    tml = {'liveboard': {'visualizations': [
        {'id': 'Viz_1', 'note_tile': {'html_parsed_string': '<div/>'}},
        {'id': 'Viz_2', 'answer': {'name': 'no query here'}},
    ]}}
    fixed_count, messages = validate_liveboard_queries(tml, column_map)
    assert fixed_count == 0
    assert messages == []