From 4d5a8b5ef7215724319fb4e9c69514b1b4631b7d Mon Sep 17 00:00:00 2001 From: Prithvi Kannan Date: Tue, 29 Oct 2024 21:40:04 -0700 Subject: [PATCH] to_markdown() in test Signed-off-by: Prithvi Kannan --- src/databricks_ai_bridge/genie.py | 4 ++-- tests/databricks_ai_bridge/test_genie.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/databricks_ai_bridge/genie.py b/src/databricks_ai_bridge/genie.py index 5a6bfb5..540d97e 100644 --- a/src/databricks_ai_bridge/genie.py +++ b/src/databricks_ai_bridge/genie.py @@ -109,8 +109,8 @@ def poll_result(): ) if resp["status"] == "EXECUTING_QUERY": query = next(r for r in resp["attachments"] if "query" in r)["query"] - description = query["description"] - sql = query["query"] + description = query.get("description", "") + sql = query.get("query", "") logging.debug(f"Description: {description}") logging.debug(f"SQL: {sql}") return poll_query_results() diff --git a/tests/databricks_ai_bridge/test_genie.py b/tests/databricks_ai_bridge/test_genie.py index edd720a..0d1cc2f 100644 --- a/tests/databricks_ai_bridge/test_genie.py +++ b/tests/databricks_ai_bridge/test_genie.py @@ -65,7 +65,7 @@ def test_poll_for_result_executing_query(genie, mock_workspace_client): }, ] result = genie.poll_for_result("123", "456") - assert result == pd.DataFrame().to_string() + assert result == pd.DataFrame().to_markdown() def test_ask_question(genie, mock_workspace_client): @@ -109,7 +109,7 @@ def test_parse_query_result_with_data(): "created_at": [datetime(2023, 10, 1).date(), datetime(2023, 10, 2).date()], } ) - assert result == expected_df.to_string() + assert result == expected_df.to_markdown() def test_parse_query_result_with_null_values(): @@ -138,4 +138,4 @@ def test_parse_query_result_with_null_values(): "created_at": [datetime(2023, 10, 1).date(), None], } ) - assert result == expected_df.to_string() + assert result == expected_df.to_markdown()