malepati/custom_template_working

1from utils.extract_response_data import extract_markdown_table_data, process_query_response2import json3 4def test_extract_markdown_table_data():5    sample_text = """6    Here is the data:7    8    | Category | Value |9    | --- | --- |10    | A | 10 |11    | B | 20 |12    | C | 30 |13    14    End of data.15    """16    data = extract_markdown_table_data(sample_text)17    assert len(data) == 318    assert data[0]['Category'] == 'A'19    assert data[0]['Value'] == 1020    print("extract_markdown_table_data passed")21 22def test_process_query_response_prioritizes_table():23    # JSON has 1 item, Table has 3 items24    response = """25    <datastory_chart>26        <vega_spec>27        {28            "data": {29                "values": [30                    {"category": "A", "value": 10}31                ]32            },33            "mark": "bar"34        }35        </vega_spec>36        37        <chart_type>bar</chart_type>38        39        | category | value |40        | --- | --- |41        | A | 10 |42        | B | 20 |43        | C | 30 |44    </datastory_chart>45    """46    47    result = process_query_response("test query", response)48    49    assert result['has_chart'] is True50    chart_data = result['chart_data']['data']51    assert len(chart_data) == 352    assert chart_data[1]['label'] == 'B'53    assert chart_data[1]['value'] == 2054    print("test_process_query_response_prioritizes_table passed")55 56if __name__ == "__main__":57    try:58        test_extract_markdown_table_data()59        test_process_query_response_prioritizes_table()60        print("All tests passed!")61    except Exception as e:62        print(f"Tests failed: {e}")63        import traceback64        traceback.print_exc()65