From 4d27cb71823340f9936de800f0d01d4721a22414 Mon Sep 17 00:00:00 2001 From: Sunish Sheth Date: Mon, 16 Dec 2024 09:51:01 -0800 Subject: [PATCH] Update based on comments Signed-off-by: Sunish Sheth --- src/databricks_ai_bridge/__init__.py | 3 --- ...ernal_tool_request.py => external_tools.py} | 2 +- src/databricks_ai_bridge/utils/annotations.py | 2 ++ ..._tool_request.py => test_external_tools.py} | 18 +++++++++++------- 4 files changed, 14 insertions(+), 11 deletions(-) rename src/databricks_ai_bridge/{external_tool_request.py => external_tools.py} (94%) rename tests/databricks_ai_bridge/{test_external_tool_request.py => test_external_tools.py} (82%) diff --git a/src/databricks_ai_bridge/__init__.py b/src/databricks_ai_bridge/__init__.py index 25165ec..e69de29 100644 --- a/src/databricks_ai_bridge/__init__.py +++ b/src/databricks_ai_bridge/__init__.py @@ -1,3 +0,0 @@ -from databricks_ai_bridge.external_tool_request import http_request - -__all__ = ["http_request"] diff --git a/src/databricks_ai_bridge/external_tool_request.py b/src/databricks_ai_bridge/external_tools.py similarity index 94% rename from src/databricks_ai_bridge/external_tool_request.py rename to src/databricks_ai_bridge/external_tools.py index b5d5144..7813b86 100644 --- a/src/databricks_ai_bridge/external_tool_request.py +++ b/src/databricks_ai_bridge/external_tools.py @@ -18,7 +18,7 @@ def http_request( params: Optional[Dict[str, Any]] = None, ) -> requests.Response: """ - Makes an HTTP request to an external function through the Databricks Workspace. + Makes an HTTP request to a remote API using authentication from a Unity Catalog HTTP connection. Args: conn (str): The connection name to use. This is required to identify the external connection. diff --git a/src/databricks_ai_bridge/utils/annotations.py b/src/databricks_ai_bridge/utils/annotations.py index 4f2b3af..098073f 100644 --- a/src/databricks_ai_bridge/utils/annotations.py +++ b/src/databricks_ai_bridge/utils/annotations.py @@ -1,3 +1,5 @@ +# This code is copied from MLflow: https://github.com/mlflow/mlflow/blob/v2.19.0/mlflow/utils/annotations.py#L31 + import inspect import re import types diff --git a/tests/databricks_ai_bridge/test_external_tool_request.py b/tests/databricks_ai_bridge/test_external_tools.py similarity index 82% rename from tests/databricks_ai_bridge/test_external_tool_request.py rename to tests/databricks_ai_bridge/test_external_tools.py index 94b9e1f..ed1d29e 100644 --- a/tests/databricks_ai_bridge/test_external_tool_request.py +++ b/tests/databricks_ai_bridge/test_external_tools.py @@ -1,15 +1,17 @@ from unittest.mock import MagicMock, patch -from databricks_ai_bridge import http_request +from databricks_ai_bridge.external_tools import http_request -@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") -@patch("databricks_ai_bridge.external_tool_request.requests.post") +@patch("databricks_ai_bridge.external_tools.WorkspaceClient") +@patch("databricks_ai_bridge.external_tools.requests.post") def test_http_request_success(mock_post, mock_workspace_client): # Mock the WorkspaceClient config mock_workspace_config = MagicMock() mock_workspace_config.host = "https://mock-host" - mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"} + mock_workspace_config._header_factory.return_value = { + "Authorization": "Bearer mock-token" + } mock_workspace_client.return_value.config = mock_workspace_config # Mock the POST request @@ -49,13 +51,15 @@ def test_http_request_success(mock_post, mock_workspace_client): ) -@patch("databricks_ai_bridge.external_tool_request.WorkspaceClient") -@patch("databricks_ai_bridge.external_tool_request.requests.post") +@patch("databricks_ai_bridge.external_tools.WorkspaceClient") +@patch("databricks_ai_bridge.external_tools.requests.post") def test_http_request_error_response(mock_post, mock_workspace_client): # Mock the WorkspaceClient config mock_workspace_config = MagicMock() mock_workspace_config.host = "https://mock-host" - mock_workspace_config._header_factory.return_value = {"Authorization": "Bearer mock-token"} + mock_workspace_config._header_factory.return_value = { + "Authorization": "Bearer mock-token" + } mock_workspace_client.return_value.config = mock_workspace_config # Mock the POST request to return an error