lazy create try

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
 
-        # 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)
+        # Store workspace client for lazy MCP client initialization
+        self._workspace_client = workspace_client
+        self._mcp_client = None  # Lazy initialization
         self._tool_name = f"query_space_{space_id}"
 
         # Keep headers for deprecated REST methods (backwards compatibility)
@@ -393,6 +393,20 @@ 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.
@@ -408,7 +422,7 @@ def ask_question(self, question, conversation_id: Optional[str] = None):
         if conversation_id:
             args["conversation_id"] = conversation_id
 
-        mcp_result = self._mcp_client.call_tool(self._tool_name, args)
+        mcp_result = self._get_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,7 +131,10 @@ def test_ask_question(genie, mock_workspace_client):
         ]
     )
 
-    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result):
+    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):
         genie_result = genie.ask_question("What is the meaning of life?")
         assert genie_result.result == "Answer"
         assert genie_result.conversation_id == "123"
@@ -147,7 +150,10 @@ def test_ask_question_continued_conversation(genie, mock_workspace_client):
         ]
     )
 
-    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result):
+    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):
         genie_result = genie.ask_question("What is the meaning of life?", "123")
         assert genie_result.result == "42"
         assert genie_result.conversation_id == "123"
@@ -163,11 +169,14 @@ def test_ask_question_calls_mcp_without_conversation_id(genie, mock_workspace_cl
         ]
     )
 
-    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result) as mock_call:
+    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):
         genie.ask_question("What is the meaning of life?")
 
         # Verify MCP client was called with correct args (no conversation_id)
-        mock_call.assert_called_once_with(
+        mock_mcp_client.call_tool.assert_called_once_with(
             "query_space_test_space_id", {"query": "What is the meaning of life?"}
         )
 
@@ -183,11 +192,14 @@ def test_ask_question_calls_mcp_with_conversation_id(genie, mock_workspace_clien
         ]
     )
 
-    with patch.object(genie._mcp_client, "call_tool", return_value=mock_mcp_result) as mock_call:
+    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):
         genie.ask_question("What is the meaning of life?", "123")
 
         # Verify MCP client was called with conversation_id included
-        mock_call.assert_called_once_with(
+        mock_mcp_client.call_tool.assert_called_once_with(
             "query_space_test_space_id",
             {"query": "What is the meaning of life?", "conversation_id": "123"},
         )