to_markdown() in test

Signed-off-by: Prithvi Kannan <[email protected]>

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()

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()