wehub-resource-sync
598 lines
18 KiB
Python
598 lines
18 KiB
Python
"""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
|