revert

Author: nisha2003

src/databricks_ai_bridge/genie.py

@@ -182,9 +182,9 @@ def __init__(
         self.genie = workspace_client.genie
         self.description = self.genie.get_space(space_id).description
 
-        # Store workspace client for lazy MCP client initialization
-        self._workspace_client = workspace_client
-        self._mcp_client = None  # Lazy initialization
+        # Initialize MCP client for communication with Genie backend
+        server_url = f"{workspace_client.config.host}/api/2.0/mcp/genie/{space_id}"
+        self._mcp_client = DatabricksMCPClient(server_url, workspace_client)
         self._tool_name = f"query_space_{space_id}"
 
         # Keep headers for deprecated REST methods (backwards compatibility)
@@ -393,20 +393,6 @@ def poll_result():
 
         return poll_result()
 
-    def _get_mcp_client(self):
-        """Lazy initialization of MCP client.
-
-        Creates the MCP client on first use to avoid issues with mocked
-        workspace clients in tests and to defer initialization until needed.
-
-        Returns:
-            DatabricksMCPClient: The initialized MCP client
-        """
-        if self._mcp_client is None:
-            server_url = f"{self._workspace_client.config.host}/api/2.0/mcp/genie/{self.space_id}"
-            self._mcp_client = DatabricksMCPClient(server_url, self._workspace_client)
-        return self._mcp_client
-
     @mlflow.trace()
     def ask_question(self, question, conversation_id: Optional[str] = None):
         """Ask a question to the Genie space using MCP protocol.
@@ -422,7 +408,7 @@ def ask_question(self, question, conversation_id: Optional[str] = None):
         if conversation_id:
             args["conversation_id"] = conversation_id
 
-        mcp_result = self._get_mcp_client().call_tool(self._tool_name, args)
+        mcp_result = self._mcp_client.call_tool(self._tool_name, args)
 
         if not mcp_result.content or len(mcp_result.content) == 0:
             return GenieResponse(

tests/databricks_ai_bridge/test_genie.py

@@ -131,10 +131,7 @@ def test_ask_question(genie, mock_workspace_client):
         ]
     )
 
-    mock_mcp_client = MagicMock()
-    mock_mcp_client.call_tool.return_value = mock_mcp_result
-
-    with patch.object(genie, "_get_mcp_client", return_value=mock_mcp_client):
+    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result):
         genie_result = genie.ask_question("What is the meaning of life?")
         assert genie_result.result == "Answer"
         assert genie_result.conversation_id == "123"
@@ -150,10 +147,7 @@ def test_ask_question_continued_conversation(genie, mock_workspace_client):
         ]
     )
 
-    mock_mcp_client = MagicMock()
-    mock_mcp_client.call_tool.return_value = mock_mcp_result
-
-    with patch.object(genie, "_get_mcp_client", return_value=mock_mcp_client):
+    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result):
         genie_result = genie.ask_question("What is the meaning of life?", "123")
         assert genie_result.result == "42"
         assert genie_result.conversation_id == "123"
@@ -169,14 +163,11 @@ def test_ask_question_calls_mcp_without_conversation_id(genie, mock_workspace_cl
         ]
     )
 
-    mock_mcp_client = MagicMock()
-    mock_mcp_client.call_tool.return_value = mock_mcp_result
-
-    with patch.object(genie, "_get_mcp_client", return_value=mock_mcp_client):
+    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result) as mock_call:
         genie.ask_question("What is the meaning of life?")
 
         # Verify MCP client was called with correct args (no conversation_id)
-        mock_mcp_client.call_tool.assert_called_once_with(
+        mock_call.assert_called_once_with(
             "query_space_test_space_id", {"query": "What is the meaning of life?"}
         )
 
@@ -192,14 +183,11 @@ def test_ask_question_calls_mcp_with_conversation_id(genie, mock_workspace_clien
         ]
     )
 
-    mock_mcp_client = MagicMock()
-    mock_mcp_client.call_tool.return_value = mock_mcp_result
-
-    with patch.object(genie, "_get_mcp_client", return_value=mock_mcp_client):
+    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result) as mock_call:
         genie.ask_question("What is the meaning of life?", "123")
 
         # Verify MCP client was called with conversation_id included
-        mock_mcp_client.call_tool.assert_called_once_with(
+        mock_call.assert_called_once_with(
             "query_space_test_space_id",
             {"query": "What is the meaning of life?", "conversation_id": "123"},
         )