chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,597 @@
|
||||
"""Unit tests for :mod:`omnigent.runner.uc_function`.
|
||||
|
||||
Tests cover:
|
||||
|
||||
- SQL statement construction with parameterized queries
|
||||
(``_build_select_statement``).
|
||||
- End-to-end ``execute_uc_function`` with a mocked
|
||||
``WorkspaceClient`` — verifies the correct SQL is sent, the
|
||||
warehouse ID is forwarded, and the result is extracted from the
|
||||
SDK response.
|
||||
- Error handling: missing warehouse ID, failed execution, no result
|
||||
data.
|
||||
- Dispatch integration: ``_is_uc_function_tool`` and
|
||||
``_execute_uc_function_tool`` in ``tool_dispatch.py``.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from omnigent.runner.uc_function import (
|
||||
_build_select_statement,
|
||||
)
|
||||
from omnigent.spec.types import LocalToolInfo, ToolRuntime
|
||||
|
||||
# ── Helpers ─────────────────────────────────────────────────────
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementStatus:
|
||||
"""Stub for ``databricks.sdk.service.sql.StatementStatus``.
|
||||
|
||||
:param state: The execution state, e.g.
|
||||
``_FakeStatementState.SUCCEEDED``.
|
||||
:param error: Optional error payload for failed statements.
|
||||
"""
|
||||
|
||||
state: Any
|
||||
error: Any = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementError:
|
||||
"""Stub for a statement execution error.
|
||||
|
||||
:param message: Human-readable error message.
|
||||
"""
|
||||
|
||||
message: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeResultData:
|
||||
"""Stub for ``databricks.sdk.service.sql.ResultData``.
|
||||
|
||||
:param data_array: The result rows, each a list of string
|
||||
column values.
|
||||
"""
|
||||
|
||||
data_array: list[list[str | None]] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementResponse:
|
||||
"""Stub for ``databricks.sdk.service.sql.StatementResponse``.
|
||||
|
||||
:param status: Execution status.
|
||||
:param result: Result data payload.
|
||||
"""
|
||||
|
||||
status: _FakeStatementStatus | None = None
|
||||
result: _FakeResultData | None = None
|
||||
|
||||
|
||||
class _FakeStatementState:
|
||||
"""Enum stand-in for ``databricks.sdk.service.sql.StatementState``.
|
||||
|
||||
Uses string ``value`` attributes matching the real enum so
|
||||
comparisons in production code work correctly.
|
||||
"""
|
||||
|
||||
SUCCEEDED = "SUCCEEDED"
|
||||
FAILED = "FAILED"
|
||||
|
||||
def __init__(self, value: str) -> None:
|
||||
self.value = value
|
||||
|
||||
def __eq__(self, other: object) -> bool:
|
||||
if isinstance(other, str):
|
||||
return self.value == other
|
||||
if hasattr(other, "value"):
|
||||
return self.value == other.value
|
||||
return NotImplemented
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeAgentSpec:
|
||||
"""Minimal agent spec stub for dispatch tests.
|
||||
|
||||
:param local_tools: List of tool info objects the spec declares.
|
||||
:param executor: Executor config stub.
|
||||
"""
|
||||
|
||||
local_tools: list[LocalToolInfo] = field(default_factory=list)
|
||||
executor: Any = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeExecutorSpec:
|
||||
"""Minimal executor spec stub.
|
||||
|
||||
:param auth: Auth config, e.g. a ``_FakeDatabricksAuth``.
|
||||
:param profile: Deprecated direct profile field.
|
||||
:param config: Compat config dict.
|
||||
"""
|
||||
|
||||
auth: Any = None
|
||||
profile: str | None = None
|
||||
config: dict[str, Any] | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FakeDatabricksAuth:
|
||||
"""Minimal Databricks auth stub.
|
||||
|
||||
:param profile: Databricks config profile name, e.g.
|
||||
``"oss"``.
|
||||
"""
|
||||
|
||||
profile: str
|
||||
|
||||
|
||||
# ── _build_select_statement tests ───────────────────────────────
|
||||
|
||||
|
||||
def test_build_select_no_args() -> None:
|
||||
"""Empty args produce a bare ``SELECT func()`` with no parameters."""
|
||||
sql, params = _build_select_statement("cat.schema.func", {})
|
||||
assert sql == "SELECT `cat.schema.func`()"
|
||||
assert params == []
|
||||
|
||||
|
||||
def test_build_select_bare_function_no_backticks() -> None:
|
||||
"""Bare function names (no dots) are emitted unquoted.
|
||||
|
||||
Built-in SQL functions like ``ai_query`` don't live in a
|
||||
catalog namespace and must not be backtick-quoted.
|
||||
"""
|
||||
sql, params = _build_select_statement("ai_query", {"endpoint": "my-ep", "prompt": "hi"})
|
||||
# No backticks around ai_query.
|
||||
assert sql == "SELECT ai_query(:endpoint, :prompt)"
|
||||
assert len(params) == 2
|
||||
|
||||
|
||||
def test_build_select_rejects_sql_injection_in_catalog_path() -> None:
|
||||
"""catalog_path with SQL metacharacters is rejected.
|
||||
|
||||
Prevents SQL injection via crafted tool specs — backticks,
|
||||
semicolons, parens, or spaces in catalog_path would alter the
|
||||
generated SQL structure.
|
||||
"""
|
||||
with pytest.raises(ValueError, match="Invalid catalog_path"):
|
||||
_build_select_statement("func(); DROP TABLE x; --", {})
|
||||
|
||||
|
||||
def test_build_select_rejects_sql_injection_in_param_name() -> None:
|
||||
"""Parameter names with SQL metacharacters are rejected.
|
||||
|
||||
Parameter names from LLM tool-call arguments are interpolated
|
||||
as ``:name`` markers. SQL metacharacters in keys would alter
|
||||
the query structure.
|
||||
"""
|
||||
with pytest.raises(ValueError, match="Invalid parameter name"):
|
||||
_build_select_statement("cat.schema.func", {"x; DROP TABLE": "val"})
|
||||
|
||||
|
||||
def test_build_select_single_arg() -> None:
|
||||
"""Single string arg produces one named parameter marker."""
|
||||
sql, params = _build_select_statement(
|
||||
"cat.schema.classify",
|
||||
{"text": "hello"},
|
||||
)
|
||||
assert sql == "SELECT `cat.schema.classify`(:text)"
|
||||
# String values are passed through directly (not JSON-encoded).
|
||||
assert len(params) == 1
|
||||
assert params[0] == {"name": "text", "value": "hello"}
|
||||
|
||||
|
||||
def test_build_select_multiple_args() -> None:
|
||||
"""Multiple args produce comma-separated parameter markers."""
|
||||
sql, params = _build_select_statement(
|
||||
"cat.schema.func",
|
||||
{"a": "x", "b": 42, "c": True},
|
||||
)
|
||||
assert ":a" in sql
|
||||
assert ":b" in sql
|
||||
assert ":c" in sql
|
||||
# Non-string values are JSON-encoded.
|
||||
b_param = next(p for p in params if p["name"] == "b")
|
||||
assert b_param["value"] == "42"
|
||||
c_param = next(p for p in params if p["name"] == "c")
|
||||
assert c_param["value"] == "true"
|
||||
|
||||
|
||||
def test_build_select_non_string_value_json_encoded() -> None:
|
||||
"""Non-string values (int, bool, list, dict) are JSON-encoded.
|
||||
|
||||
The SQL Statement Execution API accepts all parameter values
|
||||
as strings; JSON encoding preserves type information for the
|
||||
server-side parser.
|
||||
"""
|
||||
_, params = _build_select_statement(
|
||||
"cat.schema.func",
|
||||
{"items": [1, 2, 3]},
|
||||
)
|
||||
assert params[0]["value"] == json.dumps([1, 2, 3])
|
||||
|
||||
|
||||
# ── execute_uc_function tests ───────────────────────────────────
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_missing_warehouse_id() -> None:
|
||||
"""Calling without warehouse_id raises ValueError immediately.
|
||||
|
||||
This prevents a confusing SDK error downstream — the user
|
||||
gets a clear message about the missing config.
|
||||
"""
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
with pytest.raises(ValueError, match="requires a warehouse_id"):
|
||||
await execute_uc_function(
|
||||
catalog_path="cat.schema.func",
|
||||
args={"x": "1"},
|
||||
warehouse_id=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_missing_warehouse_id_env_unset(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
"""ValueError when warehouse_id is None AND env var is unset.
|
||||
|
||||
Ensures the env var fallback path is exercised: when neither
|
||||
the parameter nor ``DATABRICKS_WAREHOUSE_ID`` is available,
|
||||
the error message names both sources.
|
||||
"""
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
monkeypatch.delenv("DATABRICKS_WAREHOUSE_ID", raising=False)
|
||||
with pytest.raises(ValueError, match="requires a warehouse_id"):
|
||||
await execute_uc_function(
|
||||
catalog_path="cat.schema.func",
|
||||
args={},
|
||||
warehouse_id=None,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_warehouse_from_env(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""warehouse_id falls back to DATABRICKS_WAREHOUSE_ID env var.
|
||||
|
||||
When the caller passes ``warehouse_id=None``, the env var
|
||||
is used instead. Verifies the env value reaches the SDK call.
|
||||
"""
|
||||
fake_response = _FakeStatementResponse(
|
||||
status=_FakeStatementStatus(state=_FakeStatementState("SUCCEEDED")),
|
||||
result=_FakeResultData(data_array=[["ok"]]),
|
||||
)
|
||||
|
||||
captured_calls: list[dict[str, object]] = []
|
||||
|
||||
def _fake_execute_statement(**kwargs: object) -> _FakeStatementResponse:
|
||||
captured_calls.append(kwargs)
|
||||
return fake_response
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementExecution:
|
||||
execute_statement: object = None
|
||||
|
||||
@dataclass
|
||||
class _FakeClient:
|
||||
statement_execution: _FakeStatementExecution = field(
|
||||
default_factory=lambda: _FakeStatementExecution(
|
||||
execute_statement=_fake_execute_statement,
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"omnigent.runner.uc_function._get_workspace_client",
|
||||
lambda profile: _FakeClient(),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class _FakeParam:
|
||||
name: str
|
||||
value: str | None = None
|
||||
|
||||
monkeypatch.setattr("databricks.sdk.service.sql.StatementParameterListItem", _FakeParam)
|
||||
monkeypatch.setattr("databricks.sdk.service.sql.StatementState", _FakeStatementState)
|
||||
monkeypatch.setenv("DATABRICKS_WAREHOUSE_ID", "env-wh-456")
|
||||
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
result = await execute_uc_function(
|
||||
catalog_path="cat.schema.func",
|
||||
args={},
|
||||
warehouse_id=None,
|
||||
)
|
||||
|
||||
# Env var value reached the SDK call.
|
||||
assert result == "ok"
|
||||
assert len(captured_calls) == 1
|
||||
assert captured_calls[0]["warehouse_id"] == "env-wh-456"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_success(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Successful UC function execution returns the scalar result.
|
||||
|
||||
Verifies that:
|
||||
- The correct SQL and warehouse_id are forwarded to the SDK.
|
||||
- The scalar result from ``data_array[0][0]`` is returned.
|
||||
"""
|
||||
# Build a fake response with a scalar result.
|
||||
fake_response = _FakeStatementResponse(
|
||||
status=_FakeStatementStatus(
|
||||
state=_FakeStatementState("SUCCEEDED"),
|
||||
),
|
||||
result=_FakeResultData(data_array=[["positive"]]),
|
||||
)
|
||||
|
||||
captured_calls: list[dict[str, Any]] = []
|
||||
|
||||
def _fake_execute_statement(
|
||||
*,
|
||||
statement: str,
|
||||
warehouse_id: str,
|
||||
parameters: Any = None,
|
||||
) -> _FakeStatementResponse:
|
||||
captured_calls.append(
|
||||
{
|
||||
"statement": statement,
|
||||
"warehouse_id": warehouse_id,
|
||||
"parameters": parameters,
|
||||
}
|
||||
)
|
||||
return fake_response
|
||||
|
||||
# Stub the workspace client.
|
||||
@dataclass
|
||||
class _FakeStatementExecution:
|
||||
execute_statement: Any = None
|
||||
|
||||
@dataclass
|
||||
class _FakeClient:
|
||||
statement_execution: _FakeStatementExecution = field(
|
||||
default_factory=lambda: _FakeStatementExecution(
|
||||
execute_statement=_fake_execute_statement,
|
||||
),
|
||||
)
|
||||
|
||||
fake_client = _FakeClient()
|
||||
|
||||
monkeypatch.setattr(
|
||||
"omnigent.runner.uc_function._get_workspace_client",
|
||||
lambda profile: fake_client,
|
||||
)
|
||||
|
||||
# Also patch the SDK imports used inside execute_uc_function.
|
||||
# StatementParameterListItem is used to wrap parameters.
|
||||
@dataclass
|
||||
class _FakeParam:
|
||||
name: str
|
||||
value: str | None = None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementParameterListItem",
|
||||
_FakeParam,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementState",
|
||||
_FakeStatementState,
|
||||
)
|
||||
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
result = await execute_uc_function(
|
||||
catalog_path="cat.schema.classify",
|
||||
args={"text": "hello"},
|
||||
profile="oss",
|
||||
warehouse_id="wh-123",
|
||||
)
|
||||
|
||||
# Scalar result extracted from data_array[0][0].
|
||||
assert result == "positive"
|
||||
# Verify the SDK was called with correct warehouse.
|
||||
assert len(captured_calls) == 1
|
||||
assert captured_calls[0]["warehouse_id"] == "wh-123"
|
||||
assert "cat.schema.classify" in captured_calls[0]["statement"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_failed_statement(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""Failed statement execution raises RuntimeError with the error message."""
|
||||
fake_response = _FakeStatementResponse(
|
||||
status=_FakeStatementStatus(
|
||||
state=_FakeStatementState("FAILED"),
|
||||
error=_FakeStatementError(message="Function not found"),
|
||||
),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementExecution:
|
||||
execute_statement: Any = None
|
||||
|
||||
@dataclass
|
||||
class _FakeClient:
|
||||
statement_execution: _FakeStatementExecution = field(
|
||||
default_factory=lambda: _FakeStatementExecution(
|
||||
execute_statement=lambda **kw: fake_response,
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"omnigent.runner.uc_function._get_workspace_client",
|
||||
lambda profile: _FakeClient(),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class _FakeParam:
|
||||
name: str
|
||||
value: str | None = None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementParameterListItem",
|
||||
_FakeParam,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementState",
|
||||
_FakeStatementState,
|
||||
)
|
||||
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
with pytest.raises(RuntimeError, match="Function not found"):
|
||||
await execute_uc_function(
|
||||
catalog_path="cat.schema.missing",
|
||||
args={},
|
||||
warehouse_id="wh-123",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_execute_uc_function_null_result(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
"""NULL result from the function returns JSON null."""
|
||||
fake_response = _FakeStatementResponse(
|
||||
status=_FakeStatementStatus(
|
||||
state=_FakeStatementState("SUCCEEDED"),
|
||||
),
|
||||
result=None,
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class _FakeStatementExecution:
|
||||
execute_statement: Any = None
|
||||
|
||||
@dataclass
|
||||
class _FakeClient:
|
||||
statement_execution: _FakeStatementExecution = field(
|
||||
default_factory=lambda: _FakeStatementExecution(
|
||||
execute_statement=lambda **kw: fake_response,
|
||||
),
|
||||
)
|
||||
|
||||
monkeypatch.setattr(
|
||||
"omnigent.runner.uc_function._get_workspace_client",
|
||||
lambda profile: _FakeClient(),
|
||||
)
|
||||
|
||||
@dataclass
|
||||
class _FakeParam:
|
||||
name: str
|
||||
value: str | None = None
|
||||
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementParameterListItem",
|
||||
_FakeParam,
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
"databricks.sdk.service.sql.StatementState",
|
||||
_FakeStatementState,
|
||||
)
|
||||
|
||||
from omnigent.runner.uc_function import execute_uc_function
|
||||
|
||||
result = await execute_uc_function(
|
||||
catalog_path="cat.schema.func",
|
||||
args={},
|
||||
warehouse_id="wh-123",
|
||||
)
|
||||
|
||||
assert result == json.dumps(None)
|
||||
|
||||
|
||||
# ── Dispatch integration tests ──────────────────────────────────
|
||||
|
||||
|
||||
def test_is_uc_function_tool_true() -> None:
|
||||
"""``_is_uc_function_tool`` returns True for UC_FUNCTION tools."""
|
||||
from omnigent.runner.tool_dispatch import _is_uc_function_tool
|
||||
|
||||
spec = _FakeAgentSpec(
|
||||
local_tools=[
|
||||
LocalToolInfo(
|
||||
name="classify",
|
||||
path=None,
|
||||
language="omnigent-python-callable",
|
||||
runtime=ToolRuntime.UC_FUNCTION,
|
||||
catalog_path="cat.schema.classify",
|
||||
warehouse_id="wh-123",
|
||||
),
|
||||
],
|
||||
)
|
||||
# UC function tool is detected.
|
||||
assert _is_uc_function_tool("classify", spec) is True
|
||||
# Non-existent tool returns False.
|
||||
assert _is_uc_function_tool("nonexistent", spec) is False
|
||||
|
||||
|
||||
def test_is_uc_function_tool_false_for_server_tool() -> None:
|
||||
"""``_is_uc_function_tool`` returns False for SERVER-runtime tools."""
|
||||
from omnigent.runner.tool_dispatch import _is_uc_function_tool
|
||||
|
||||
spec = _FakeAgentSpec(
|
||||
local_tools=[
|
||||
LocalToolInfo(
|
||||
name="calc",
|
||||
path="tests.tool_functions.calculate",
|
||||
language="python",
|
||||
runtime=ToolRuntime.SERVER,
|
||||
),
|
||||
],
|
||||
)
|
||||
assert _is_uc_function_tool("calc", spec) is False
|
||||
|
||||
|
||||
def test_is_uc_function_tool_no_spec() -> None:
|
||||
"""``_is_uc_function_tool`` returns False when agent_spec is None."""
|
||||
from omnigent.runner.tool_dispatch import _is_uc_function_tool
|
||||
|
||||
assert _is_uc_function_tool("classify", None) is False
|
||||
|
||||
|
||||
def test_resolve_uc_profile_from_auth() -> None:
|
||||
"""Profile is extracted from ``executor.auth.profile``."""
|
||||
from omnigent.runner.tool_dispatch import _resolve_uc_profile
|
||||
|
||||
spec = _FakeAgentSpec(
|
||||
executor=_FakeExecutorSpec(
|
||||
auth=_FakeDatabricksAuth(profile="oss"),
|
||||
),
|
||||
)
|
||||
assert _resolve_uc_profile(spec) == "oss"
|
||||
|
||||
|
||||
def test_resolve_uc_profile_from_deprecated_field() -> None:
|
||||
"""Profile falls back to ``executor.profile`` when auth is absent."""
|
||||
from omnigent.runner.tool_dispatch import _resolve_uc_profile
|
||||
|
||||
spec = _FakeAgentSpec(
|
||||
executor=_FakeExecutorSpec(profile="legacy-profile"),
|
||||
)
|
||||
assert _resolve_uc_profile(spec) == "legacy-profile"
|
||||
|
||||
|
||||
def test_resolve_uc_profile_from_config() -> None:
|
||||
"""Profile falls back to ``executor.config["profile"]``."""
|
||||
from omnigent.runner.tool_dispatch import _resolve_uc_profile
|
||||
|
||||
spec = _FakeAgentSpec(
|
||||
executor=_FakeExecutorSpec(config={"profile": "compat-profile"}),
|
||||
)
|
||||
assert _resolve_uc_profile(spec) == "compat-profile"
|
||||
|
||||
|
||||
def test_resolve_uc_profile_none() -> None:
|
||||
"""Returns None when no profile is configured anywhere."""
|
||||
from omnigent.runner.tool_dispatch import _resolve_uc_profile
|
||||
|
||||
spec = _FakeAgentSpec(executor=_FakeExecutorSpec())
|
||||
assert _resolve_uc_profile(spec) is None
|
||||
Reference in New Issue
Block a user