nvidia-nat-test 1.3.0.dev2__py3-none-any.whl → 1.3.0rc2__py3-none-any.whl

nat/test/functions.py

@@ -21,6 +21,7 @@ from nat.cli.register_workflow import register_function
 from nat.data_models.api_server import ChatRequest
 from nat.data_models.api_server import ChatResponse
 from nat.data_models.api_server import ChatResponseChunk
+from nat.data_models.api_server import Usage
 from nat.data_models.function import FunctionBaseConfig
 
 
@@ -35,7 +36,14 @@ async def echo_function(config: EchoFunctionConfig, builder: Builder):
         return message
 
     async def inner_oai(message: ChatRequest) -> ChatResponse:
-        return ChatResponse.from_string(message.messages[0].content)
+        content = message.messages[0].content
+
+        # Create usage statistics for the response
+        prompt_tokens = sum(len(str(msg.content).split()) for msg in message.messages)
+        completion_tokens = len(content.split()) if content else 0
+        total_tokens = prompt_tokens + completion_tokens
+        usage = Usage(prompt_tokens=prompt_tokens, completion_tokens=completion_tokens, total_tokens=total_tokens)
+        return ChatResponse.from_string(content, usage=usage)
 
     if (config.use_openai_api):
         yield inner_oai

nat/test/llm.py

@@ -0,0 +1,205 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# pylint: disable=unused-argument,missing-class-docstring,missing-function-docstring,import-outside-toplevel
+# pylint: disable=too-few-public-methods
+
+import asyncio
+import time
+from collections.abc import AsyncGenerator
+from collections.abc import Iterator
+from itertools import cycle as iter_cycle
+from typing import Any
+
+from pydantic import Field
+
+from nat.builder.builder import Builder
+from nat.builder.framework_enum import LLMFrameworkEnum
+from nat.builder.llm import LLMProviderInfo
+from nat.cli.register_workflow import register_llm_client
+from nat.cli.register_workflow import register_llm_provider
+from nat.data_models.llm import LLMBaseConfig
+
+
+class TestLLMConfig(LLMBaseConfig, name="nat_test_llm"):
+    """Test LLM configuration."""
+    __test__ = False
+    response_seq: list[str] = Field(
+        default=[],
+        description="Returns the next element in order (wraps)",
+    )
+    delay_ms: int = Field(default=0, ge=0, description="Artificial per-call delay in milliseconds to mimic latency")
+
+
+class _ResponseChooser:
+    """
+    Helper class to choose the next response according to config using itertools.cycle and provide synchronous and
+    asynchronous sleep functions.
+    """
+
+    def __init__(self, response_seq: list[str], delay_ms: int):
+        self._cycler = iter_cycle(response_seq) if response_seq else None
+        self._delay_ms = delay_ms
+
+    def next_response(self) -> str:
+        """Return the next response in the cycle, or an empty string if no responses are configured."""
+        if self._cycler is None:
+            return ""
+        return next(self._cycler)
+
+    def sync_sleep(self) -> None:
+        time.sleep(self._delay_ms / 1000.0)
+
+    async def async_sleep(self) -> None:
+        await asyncio.sleep(self._delay_ms / 1000.0)
+
+
+@register_llm_provider(config_type=TestLLMConfig)
+async def test_llm_provider(config: TestLLMConfig, builder: Builder):
+    """Register the `nat_test_llm` provider for the NAT registry."""
+    yield LLMProviderInfo(config=config, description="Test LLM provider")
+
+
+@register_llm_client(config_type=TestLLMConfig, wrapper_type=LLMFrameworkEnum.LANGCHAIN)
+async def test_llm_langchain(config: TestLLMConfig, builder: Builder):
+    """LLM client for LangChain/LangGraph."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class LangChainTestLLM:
+
+        def invoke(self, messages: Any, **_kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+        async def ainvoke(self, messages: Any, **_kwargs: Any) -> str:
+            await chooser.async_sleep()
+            return chooser.next_response()
+
+        def stream(self, messages: Any, **_kwargs: Any) -> Iterator[str]:
+            chooser.sync_sleep()
+            yield chooser.next_response()
+
+        async def astream(self, messages: Any, **_kwargs: Any) -> AsyncGenerator[str]:
+            await chooser.async_sleep()
+            yield chooser.next_response()
+
+    yield LangChainTestLLM()
+
+
+@register_llm_client(config_type=TestLLMConfig, wrapper_type=LLMFrameworkEnum.LLAMA_INDEX)
+async def test_llm_llama_index(config: TestLLMConfig, builder: Builder):
+
+    try:
+        from llama_index.core.base.llms.types import ChatMessage
+        from llama_index.core.base.llms.types import ChatResponse
+    except ImportError as exc:
+        raise ImportError("llama_index is required for using the test_llm with llama_index. "
+                          "Please install the `nvidia-nat-llama-index` package. ") from exc
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class LITestLLM:
+
+        def chat(self, messages: list[Any] | None = None, **_kwargs: Any) -> ChatResponse:
+            chooser.sync_sleep()
+            return ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        async def achat(self, messages: list[Any] | None = None, **_kwargs: Any) -> ChatResponse:
+            await chooser.async_sleep()
+            return ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        def stream_chat(self, messages: list[Any] | None = None, **_kwargs: Any) -> Iterator[ChatResponse]:
+            chooser.sync_sleep()
+            yield ChatResponse(message=ChatMessage(chooser.next_response()))
+
+        async def astream_chat(self,
+                               messages: list[Any] | None = None,
+                               **_kwargs: Any) -> AsyncGenerator[ChatResponse, None]:
+            await chooser.async_sleep()
+            yield ChatResponse(message=ChatMessage(chooser.next_response()))
+
+    yield LITestLLM()
+
+
+@register_llm_client(config_type=TestLLMConfig, wrapper_type=LLMFrameworkEnum.CREWAI)
+async def test_llm_crewai(config: TestLLMConfig, builder: Builder):
+    """LLM client for CrewAI."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class CrewAITestLLM:
+
+        def call(self, messages: list[dict[str, str]] | None = None, **kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+    yield CrewAITestLLM()
+
+
+@register_llm_client(config_type=TestLLMConfig, wrapper_type=LLMFrameworkEnum.SEMANTIC_KERNEL)
+async def test_llm_semantic_kernel(config: TestLLMConfig, builder: Builder):
+    """LLM client for SemanticKernel."""
+
+    try:
+        from semantic_kernel.contents.chat_message_content import ChatMessageContent
+        from semantic_kernel.contents.utils.author_role import AuthorRole
+    except ImportError as exc:
+        raise ImportError("Semantic Kernel is required for using the test_llm with semantic_kernel. "
+                          "Please install the `nvidia-nat-semantic-kernel` package. ") from exc
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class SKTestLLM:
+
+        async def get_chat_message_contents(self, chat_history: Any, **_kwargs: Any) -> list[ChatMessageContent]:
+            await chooser.async_sleep()
+            text = chooser.next_response()
+            return [ChatMessageContent(role=AuthorRole.ASSISTANT, content=text)]
+
+        async def get_streaming_chat_message_contents(self, chat_history: Any,
+                                                      **_kwargs: Any) -> AsyncGenerator[ChatMessageContent, None]:
+            await chooser.async_sleep()
+            text = chooser.next_response()
+            yield ChatMessageContent(role=AuthorRole.ASSISTANT, content=text)
+
+    yield SKTestLLM()
+
+
+@register_llm_client(config_type=TestLLMConfig, wrapper_type=LLMFrameworkEnum.AGNO)
+async def test_llm_agno(config: TestLLMConfig, builder: Builder):
+    """LLM client for agno."""
+
+    chooser = _ResponseChooser(response_seq=config.response_seq, delay_ms=config.delay_ms)
+
+    class AgnoTestLLM:
+
+        def invoke(self, messages: Any | None = None, **_kwargs: Any) -> str:
+            chooser.sync_sleep()
+            return chooser.next_response()
+
+        async def ainvoke(self, messages: Any | None = None, **_kwargs: Any) -> str:
+            await chooser.async_sleep()
+            return chooser.next_response()
+
+        def invoke_stream(self, messages: Any | None = None, **_kwargs: Any) -> Iterator[str]:
+            chooser.sync_sleep()
+            yield chooser.next_response()
+
+        async def ainvoke_stream(self, messages: Any | None = None, **_kwargs: Any) -> AsyncGenerator[str, None]:
+            await chooser.async_sleep()
+            yield chooser.next_response()
+
+    yield AgnoTestLLM()

nat/test/plugin.py

@@ -13,20 +13,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import os
+import subprocess
+import typing
+from pathlib import Path
+
 import pytest
 
+if typing.TYPE_CHECKING:
+    from docker.client import DockerClient
+
 
 def pytest_addoption(parser: pytest.Parser):
     """
     Adds command line options for running specfic tests that are disabled by default
     """
-    parser.addoption(
-        "--run_e2e",
-        action="store_true",
-        dest="run_e2e",
-        help="Run end to end tests that would otherwise be skipped",
-    )
-
     parser.addoption(
         "--run_integration",
         action="store_true",
@@ -42,12 +43,16 @@ def pytest_addoption(parser: pytest.Parser):
         help="Run end to end tests that would otherwise be skipped",
     )
 
+    parser.addoption(
+        "--fail_missing",
+        action="store_true",
+        dest="fail_missing",
+        help=("Tests requiring unmet dependencies are normally skipped. "
+              "Setting this flag will instead cause them to be reported as a failure"),
+    )
+
 
 def pytest_runtest_setup(item):
-    if (not item.config.getoption("--run_e2e")):
-        if (item.get_closest_marker("e2e") is not None):
-            pytest.skip("Skipping end to end tests by default. Use --run_e2e to enable")
-
     if (not item.config.getoption("--run_integration")):
         if (item.get_closest_marker("integration") is not None):
             pytest.skip("Skipping integration tests by default. Use --run_integration to enable")
@@ -68,7 +73,6 @@ def register_components_fixture():
     discover_and_register_plugins(PluginTypes.ALL)
 
     # Also import the nat.test.register module to register test-only components
-    import nat.test.register  # pylint: disable=unused-import # noqa: F401
 
 
 @pytest.fixture(name="module_registry", scope="module", autouse=True)
@@ -95,3 +99,236 @@ def function_registry_fixture():
 
     with GlobalTypeRegistry.push() as registry:
         yield registry
+
+
+@pytest.fixture(scope="session", name="fail_missing")
+def fail_missing_fixture(pytestconfig: pytest.Config) -> bool:
+    """
+    Returns the value of the `fail_missing` flag, when false tests requiring unmet dependencies will be skipped, when
+    True they will fail.
+    """
+    yield pytestconfig.getoption("fail_missing")
+
+
+def require_env_variables(varnames: list[str], reason: str, fail_missing: bool = False) -> dict[str, str]:
+    """
+    Checks if the given environment variable is set, and returns its value if it is. If the variable is not set, and
+    `fail_missing` is False the test will ve skipped, otherwise a `RuntimeError` will be raised.
+    """
+    env_variables = {}
+    try:
+        for varname in varnames:
+            env_variables[varname] = os.environ[varname]
+    except KeyError as e:
+        if fail_missing:
+            raise RuntimeError(reason) from e
+
+        pytest.skip(reason=reason)
+
+    return env_variables
+
+
+@pytest.fixture(name="openai_api_key", scope='session')
+def openai_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require an Openai API key.
+    """
+    yield require_env_variables(
+        varnames=["OPENAI_API_KEY"],
+        reason="openai integration tests require the `OPENAI_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="nvidia_api_key", scope='session')
+def nvidia_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require an Nvidia API key.
+    """
+    yield require_env_variables(
+        varnames=["NVIDIA_API_KEY"],
+        reason="Nvidia integration tests require the `NVIDIA_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="serp_api_key", scope='session')
+def serp_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require a SERP API key.
+    """
+    yield require_env_variables(
+        varnames=["SERP_API_KEY"],
+        reason="SERP integration tests require the `SERP_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="tavily_api_key", scope='session')
+def tavily_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require a Tavily API key.
+    """
+    yield require_env_variables(
+        varnames=["TAVILY_API_KEY"],
+        reason="Tavily integration tests require the `TAVILY_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="mem0_api_key", scope='session')
+def mem0_api_key_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require a Mem0 API key.
+    """
+    yield require_env_variables(
+        varnames=["MEM0_API_KEY"],
+        reason="Mem0 integration tests require the `MEM0_API_KEY` environment variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="aws_keys", scope='session')
+def aws_keys_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require AWS credentials.
+    """
+
+    yield require_env_variables(
+        varnames=["AWS_ACCESS_KEY_ID", "AWS_SECRET_ACCESS_KEY"],
+        reason=
+        "AWS integration tests require the `AWS_ACCESS_KEY_ID` and `AWS_SECRET_ACCESS_KEY` environment variables to be "
+        "defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="azure_openai_keys", scope='session')
+def azure_openai_keys_fixture(fail_missing: bool):
+    """
+    Use for integration tests that require Azure OpenAI credentials.
+    """
+    yield require_env_variables(
+        varnames=["AZURE_OPENAI_API_KEY", "AZURE_OPENAI_ENDPOINT"],
+        reason="Azure integration tests require the `AZURE_OPENAI_API_KEY` and `AZURE_OPENAI_ENDPOINT` environment "
+        "variable to be defined.",
+        fail_missing=fail_missing)
+
+
+@pytest.fixture(name="require_docker", scope='session')
+def require_docker_fixture(fail_missing: bool) -> "DockerClient":
+    """
+    Use for integration tests that require Docker to be running.
+    """
+    try:
+        from docker.client import DockerClient
+        yield DockerClient()
+    except Exception as e:
+        reason = f"Unable to connect to Docker daemon: {e}"
+        if fail_missing:
+            raise RuntimeError(reason) from e
+        pytest.skip(reason=reason)
+
+
+@pytest.fixture(name="restore_environ")
+def restore_environ_fixture():
+    orig_vars = os.environ.copy()
+    yield os.environ
+
+    for key, value in orig_vars.items():
+        os.environ[key] = value
+
+    # Delete any new environment variables
+    # Iterating over a copy of the keys as we will potentially be deleting keys in the loop
+    for key in list(os.environ.keys()):
+        if key not in orig_vars:
+            del (os.environ[key])
+
+
+@pytest.fixture(name="root_repo_dir", scope='session')
+def root_repo_dir_fixture() -> Path:
+    from nat.test.utils import locate_repo_root
+    return locate_repo_root()
+
+
+@pytest.fixture(name="require_etcd", scope="session")
+def require_etcd_fixture(fail_missing: bool = False) -> bool:
+    """
+    To run these tests, an etcd server must be running
+    """
+    import requests
+
+    host = os.getenv("NAT_CI_ETCD_HOST", "localhost")
+    port = os.getenv("NAT_CI_ETCD_PORT", "2379")
+    health_url = f"http://{host}:{port}/health"
+
+    try:
+        response = requests.get(health_url, timeout=5)
+        response.raise_for_status()
+        return True
+    except:  # noqa: E722
+        failure_reason = f"Unable to connect to etcd server at {health_url}"
+        if fail_missing:
+            raise RuntimeError(failure_reason)
+        pytest.skip(reason=failure_reason)
+
+
+@pytest.fixture(name="milvus_uri", scope="session")
+def milvus_uri_fixture(require_etcd: bool, fail_missing: bool = False) -> str:
+    """
+    To run these tests, a Milvus server must be running
+    """
+    host = os.getenv("NAT_CI_MILVUS_HOST", "localhost")
+    port = os.getenv("NAT_CI_MILVUS_PORT", "19530")
+    uri = f"http://{host}:{port}"
+    try:
+        from pymilvus import MilvusClient
+        MilvusClient(uri=uri)
+
+        return uri
+    except:  # noqa: E722
+        reason = f"Unable to connect to Milvus server at {uri}"
+        if fail_missing:
+            raise RuntimeError(reason)
+        pytest.skip(reason=reason)
+
+
+@pytest.fixture(name="populate_milvus", scope="session")
+def populate_milvus_fixture(milvus_uri: str, root_repo_dir: Path):
+    """
+    Populate Milvus with some test data.
+    """
+    populate_script = root_repo_dir / "scripts/langchain_web_ingest.py"
+
+    # Ingest default cuda docs
+    subprocess.run(["python", str(populate_script), "--milvus_uri", milvus_uri], check=True)
+
+    # Ingest MCP docs
+    subprocess.run([
+        "python",
+        str(populate_script),
+        "--milvus_uri",
+        milvus_uri,
+        "--urls",
+        "https://github.com/modelcontextprotocol/python-sdk",
+        "--urls",
+        "https://modelcontextprotocol.io/introduction",
+        "--urls",
+        "https://modelcontextprotocol.io/quickstart/server",
+        "--urls",
+        "https://modelcontextprotocol.io/quickstart/client",
+        "--urls",
+        "https://modelcontextprotocol.io/examples",
+        "--urls",
+        "https://modelcontextprotocol.io/docs/concepts/architecture",
+        "--collection_name",
+        "mcp_docs"
+    ],
+                   check=True)
+
+    # Ingest some wikipedia docs
+    subprocess.run([
+        "python",
+        str(populate_script),
+        "--milvus_uri",
+        milvus_uri,
+        "--urls",
+        "https://en.wikipedia.org/wiki/Aardvark",
+        "--collection_name",
+        "wikipedia_docs"
+    ],
+                   check=True)

nat/test/register.py

@@ -13,7 +13,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# pylint: disable=unused-import
 # flake8: noqa
 # isort:skip_file
 
@@ -22,3 +21,5 @@
 from . import embedder
 from . import functions
 from . import memory
+from . import llm
+from . import utils

nat/test/tool_test_runner.py

@@ -14,18 +14,33 @@
 # limitations under the License.
 
 import asyncio
+import inspect
 import logging
 import typing
+from collections.abc import Sequence
 from contextlib import asynccontextmanager
 from unittest.mock import AsyncMock
 from unittest.mock import MagicMock
 
+from nat.authentication.interfaces import AuthProviderBase
 from nat.builder.builder import Builder
 from nat.builder.function import Function
+from nat.builder.function import FunctionGroup
 from nat.builder.function_info import FunctionInfo
 from nat.cli.type_registry import GlobalTypeRegistry
+from nat.data_models.authentication import AuthProviderBaseConfig
+from nat.data_models.embedder import EmbedderBaseConfig
 from nat.data_models.function import FunctionBaseConfig
+from nat.data_models.function import FunctionGroupBaseConfig
+from nat.data_models.function_dependencies import FunctionDependencies
+from nat.data_models.llm import LLMBaseConfig
+from nat.data_models.memory import MemoryBaseConfig
 from nat.data_models.object_store import ObjectStoreBaseConfig
+from nat.data_models.retriever import RetrieverBaseConfig
+from nat.data_models.ttc_strategy import TTCStrategyBaseConfig
+from nat.experimental.test_time_compute.models.stage_enums import PipelineTypeEnum
+from nat.experimental.test_time_compute.models.stage_enums import StageTypeEnum
+from nat.memory.interfaces import MemoryEditor
 from nat.object_store.interfaces import ObjectStore
 from nat.runtime.loader import PluginTypes
 from nat.runtime.loader import discover_and_register_plugins
@@ -46,6 +61,10 @@ class MockBuilder(Builder):
         """Add a mock function that returns a fixed response."""
         self._mocks[name] = mock_response
 
+    def mock_function_group(self, name: str, mock_response: typing.Any):
+        """Add a mock function group that returns a fixed response."""
+        self._mocks[name] = mock_response
+
     def mock_llm(self, name: str, mock_response: typing.Any):
         """Add a mock LLM that returns a fixed response."""
         self._mocks[f"llm_{name}"] = mock_response
@@ -70,14 +89,16 @@ class MockBuilder(Builder):
         """Add a mock TTC strategy that returns a fixed response."""
         self._mocks[f"ttc_strategy_{name}"] = mock_response
 
-    async def add_ttc_strategy(self, name: str, config):
+    def mock_auth_provider(self, name: str, mock_response: typing.Any):
+        """Add a mock auth provider that returns a fixed response."""
+        self._mocks[f"auth_provider_{name}"] = mock_response
+
+    async def add_ttc_strategy(self, name: str, config: TTCStrategyBaseConfig) -> None:
         """Mock implementation (no‑op)."""
         pass
 
-    async def get_ttc_strategy(self,
-                               strategy_name: str,
-                               pipeline_type: typing.Any = None,
-                               stage_type: typing.Any = None):
+    async def get_ttc_strategy(self, strategy_name: str, pipeline_type: PipelineTypeEnum,
+                               stage_type: StageTypeEnum) -> typing.Any:
         """Return a mock TTC strategy if one is configured."""
         key = f"ttc_strategy_{strategy_name}"
         if key in self._mocks:
@@ -90,16 +111,29 @@ class MockBuilder(Builder):
 
     async def get_ttc_strategy_config(self,
                                       strategy_name: str,
-                                      pipeline_type: typing.Any = None,
-                                      stage_type: typing.Any = None):
+                                      pipeline_type: PipelineTypeEnum,
+                                      stage_type: StageTypeEnum) -> TTCStrategyBaseConfig:
         """Mock implementation."""
+        return TTCStrategyBaseConfig()
+
+    async def add_auth_provider(self, name: str, config: AuthProviderBaseConfig) -> None:
+        """Mock implementation (no‑op)."""
         pass
 
+    async def get_auth_provider(self, auth_provider_name: str) -> AuthProviderBase:
+        """Return a mock auth provider if one is configured."""
+        key = f"auth_provider_{auth_provider_name}"
+        if key in self._mocks:
+            mock_auth = MagicMock()
+            mock_auth.authenticate = AsyncMock(return_value=self._mocks[key])
+            return mock_auth
+        raise ValueError(f"Auth provider '{auth_provider_name}' not mocked. Use mock_auth_provider() to add it.")
+
     async def add_function(self, name: str, config: FunctionBaseConfig) -> Function:
         """Mock implementation - not used in tool testing."""
         raise NotImplementedError("Mock implementation does not support add_function")
 
-    def get_function(self, name: str) -> Function:
+    async def get_function(self, name: str) -> Function:
         """Return a mock function if one is configured."""
         if name in self._mocks:
             mock_fn = AsyncMock()
@@ -109,25 +143,49 @@ class MockBuilder(Builder):
 
     def get_function_config(self, name: str) -> FunctionBaseConfig:
         """Mock implementation."""
-        pass
+        return FunctionBaseConfig()
+
+    async def add_function_group(self, name: str, config: FunctionGroupBaseConfig) -> FunctionGroup:
+        """Mock implementation - not used in tool testing."""
+        raise NotImplementedError("Mock implementation does not support add_function_group")
+
+    async def get_function_group(self, name: str) -> FunctionGroup:
+        """Return a mock function group if one is configured."""
+        if name in self._mocks:
+            mock_fn_group = MagicMock(spec=FunctionGroup)
+            mock_fn_group.ainvoke = AsyncMock(return_value=self._mocks[name])
+            return mock_fn_group
+        raise ValueError(f"Function group '{name}' not mocked. Use mock_function_group() to add it.")
+
+    def get_function_group_config(self, name: str) -> FunctionGroupBaseConfig:
+        """Mock implementation."""
+        return FunctionGroupBaseConfig()
 
     async def set_workflow(self, config: FunctionBaseConfig) -> Function:
         """Mock implementation."""
-        pass
+        mock_fn = AsyncMock()
+        mock_fn.ainvoke = AsyncMock(return_value="mock_workflow_result")
+        return mock_fn
 
     def get_workflow(self) -> Function:
         """Mock implementation."""
-        pass
+        mock_fn = AsyncMock()
+        mock_fn.ainvoke = AsyncMock(return_value="mock_workflow_result")
+        return mock_fn
 
     def get_workflow_config(self) -> FunctionBaseConfig:
         """Mock implementation."""
-        pass
+        return FunctionBaseConfig()
+
+    async def get_tools(self, tool_names: Sequence[str], wrapper_type) -> list[typing.Any]:
+        """Mock implementation."""
+        return []
 
-    def get_tool(self, fn_name: str, wrapper_type):
+    async def get_tool(self, fn_name: str, wrapper_type) -> typing.Any:
         """Mock implementation."""
         pass
 
-    async def add_llm(self, name: str, config):
+    async def add_llm(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -141,11 +199,11 @@ class MockBuilder(Builder):
             return mock_llm
         raise ValueError(f"LLM '{llm_name}' not mocked. Use mock_llm() to add it.")
 
-    def get_llm_config(self, llm_name: str):
+    def get_llm_config(self, llm_name: str) -> LLMBaseConfig:
         """Mock implementation."""
-        pass
+        return LLMBaseConfig()
 
-    async def add_embedder(self, name: str, config):
+    async def add_embedder(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -159,15 +217,14 @@ class MockBuilder(Builder):
             return mock_embedder
         raise ValueError(f"Embedder '{embedder_name}' not mocked. Use mock_embedder() to add it.")
 
-    def get_embedder_config(self, embedder_name: str):
+    def get_embedder_config(self, embedder_name: str) -> EmbedderBaseConfig:
         """Mock implementation."""
-        pass
+        return EmbedderBaseConfig()
 
-    async def add_memory_client(self, name: str, config):
-        """Mock implementation."""
-        pass
+    async def add_memory_client(self, name: str, config) -> MemoryEditor:
+        return MagicMock(spec=MemoryEditor)
 
-    def get_memory_client(self, memory_name: str):
+    async def get_memory_client(self, memory_name: str) -> MemoryEditor:
         """Return a mock memory client if one is configured."""
         key = f"memory_{memory_name}"
         if key in self._mocks:
@@ -177,11 +234,11 @@ class MockBuilder(Builder):
             return mock_memory
         raise ValueError(f"Memory client '{memory_name}' not mocked. Use mock_memory_client() to add it.")
 
-    def get_memory_client_config(self, memory_name: str):
+    def get_memory_client_config(self, memory_name: str) -> MemoryBaseConfig:
         """Mock implementation."""
-        pass
+        return MemoryBaseConfig()
 
-    async def add_retriever(self, name: str, config):
+    async def add_retriever(self, name: str, config) -> None:
         """Mock implementation."""
         pass
 
@@ -194,13 +251,13 @@ class MockBuilder(Builder):
             return mock_retriever
         raise ValueError(f"Retriever '{retriever_name}' not mocked. Use mock_retriever() to add it.")
 
-    async def get_retriever_config(self, retriever_name: str):
+    async def get_retriever_config(self, retriever_name: str) -> RetrieverBaseConfig:
         """Mock implementation."""
-        pass
+        return RetrieverBaseConfig()
 
-    async def add_object_store(self, name: str, config: ObjectStoreBaseConfig):
+    async def add_object_store(self, name: str, config: ObjectStoreBaseConfig) -> ObjectStore:
         """Mock implementation for object store."""
-        pass
+        return MagicMock(spec=ObjectStore)
 
     async def get_object_store_client(self, object_store_name: str) -> ObjectStore:
         """Return a mock object store client if one is configured."""
@@ -216,7 +273,7 @@ class MockBuilder(Builder):
 
     def get_object_store_config(self, object_store_name: str) -> ObjectStoreBaseConfig:
         """Mock implementation for object store config."""
-        pass
+        return ObjectStoreBaseConfig()
 
     def get_user_manager(self):
         """Mock implementation."""
@@ -224,9 +281,13 @@ class MockBuilder(Builder):
         mock_user.get_id = MagicMock(return_value="test_user")
         return mock_user
 
-    def get_function_dependencies(self, fn_name: str):
+    def get_function_dependencies(self, fn_name: str) -> FunctionDependencies:
         """Mock implementation."""
-        pass
+        return FunctionDependencies()
+
+    def get_function_group_dependencies(self, fn_name: str) -> FunctionDependencies:
+        """Mock implementation."""
+        return FunctionDependencies()
 
 
 class ToolTestRunner:
@@ -322,15 +383,19 @@ class ToolTestRunner:
 
             # Execute the tool
             if input_data is not None:
-                if asyncio.iscoroutinefunction(tool_function):
+                if isinstance(tool_function, Function):
+                    result = await tool_function.ainvoke(input_data)
+                elif asyncio.iscoroutinefunction(tool_function):
                     result = await tool_function(input_data)
                 else:
                     result = tool_function(input_data)
+            elif isinstance(tool_function, Function):
+                # Function objects require input, so pass None if no input_data
+                result = await tool_function.ainvoke(None)
+            elif asyncio.iscoroutinefunction(tool_function):
+                result = await tool_function()
             else:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function()
-                else:
-                    result = tool_function()
+                result = tool_function()
 
             # Assert expected output if provided
             if expected_output is not None:
@@ -402,8 +467,8 @@ class ToolTestRunner:
             elif isinstance(tool_result, FunctionInfo):
                 if tool_result.single_fn:
                     tool_function = tool_result.single_fn
-                elif tool_result.streaming_fn:
-                    tool_function = tool_result.streaming_fn
+                elif tool_result.stream_fn:
+                    tool_function = tool_result.stream_fn
                 else:
                     raise ValueError("Tool function not found in FunctionInfo")
             elif callable(tool_result):
@@ -413,15 +478,17 @@ class ToolTestRunner:
 
             # Execute the tool
             if input_data is not None:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function(input_data)
+                if isinstance(tool_function, Function):
+                    result = await tool_function.ainvoke(input_data)
                 else:
-                    result = tool_function(input_data)
+                    maybe_result = tool_function(input_data)
+                    result = await maybe_result if inspect.isawaitable(maybe_result) else maybe_result
+            elif isinstance(tool_function, Function):
+                # Function objects require input, so pass None if no input_data
+                result = await tool_function.ainvoke(None)
             else:
-                if asyncio.iscoroutinefunction(tool_function):
-                    result = await tool_function()
-                else:
-                    result = tool_function()
+                maybe_result = tool_function()
+                result = await maybe_result if inspect.isawaitable(maybe_result) else maybe_result
 
             # Assert expected output if provided
             if expected_output is not None:

nat/test/utils.py

@@ -0,0 +1,87 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import importlib.resources
+import inspect
+import subprocess
+import typing
+from pathlib import Path
+
+if typing.TYPE_CHECKING:
+    from nat.data_models.config import Config
+    from nat.utils.type_utils import StrPath
+
+
+def locate_repo_root() -> Path:
+    result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
+    assert result.returncode == 0, f"Failed to get git root: {result.stderr}"
+    return Path(result.stdout.strip())
+
+
+def locate_example_src_dir(example_config_class: type) -> Path:
+    """
+    Locate the example src directory for an example's config class.
+    """
+    package_name = inspect.getmodule(example_config_class).__package__
+    return importlib.resources.files(package_name)
+
+
+def locate_example_dir(example_config_class: type) -> Path:
+    """
+    Locate the example directory for an example's config class.
+    """
+    src_dir = locate_example_src_dir(example_config_class)
+    example_dir = src_dir.parent.parent
+    return example_dir
+
+
+def locate_example_config(example_config_class: type,
+                          config_file: str = "config.yml",
+                          assert_exists: bool = True) -> Path:
+    """
+    Locate the example config file for an example's config class, assumes the example contains a 'configs' directory
+    """
+    example_dir = locate_example_src_dir(example_config_class)
+    config_path = example_dir.joinpath("configs", config_file).absolute()
+    if assert_exists:
+        assert config_path.exists(), f"Config file {config_path} does not exist"
+
+    return config_path
+
+
+async def run_workflow(
+    config_file: "StrPath | None",
+    question: str,
+    expected_answer: str,
+    assert_expected_answer: bool = True,
+    config: "Config | None" = None,
+) -> str:
+    from nat.builder.workflow_builder import WorkflowBuilder
+    from nat.runtime.loader import load_config
+    from nat.runtime.session import SessionManager
+
+    if config is None:
+        assert config_file is not None, "Either config_file or config must be provided"
+        config = load_config(config_file)
+
+    async with WorkflowBuilder.from_config(config=config) as workflow_builder:
+        workflow = SessionManager(await workflow_builder.build())
+        async with workflow.run(question) as runner:
+            result = await runner.result(to_type=str)
+
+    if assert_expected_answer:
+        assert expected_answer.lower() in result.lower(), f"Expected '{expected_answer}' in '{result}'"
+
+    return result

{nvidia_nat_test-1.3.0.dev2.dist-info → nvidia_nat_test-1.3.0rc2.dist-info}/METADATA

@@ -1,12 +1,15 @@
 Metadata-Version: 2.4
 Name: nvidia-nat-test
-Version: 1.3.0.dev2
+Version: 1.3.0rc2
 Summary: Testing utilities for NeMo Agent toolkit
 Keywords: ai,rag,agents
 Classifier: Programming Language :: Python
-Requires-Python: <3.13,>=3.11
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Requires-Python: <3.14,>=3.11
 Description-Content-Type: text/markdown
-Requires-Dist: nvidia-nat==v1.3.0-dev2
+Requires-Dist: nvidia-nat==v1.3.0-rc2
 Requires-Dist: langchain-community~=0.3
 Requires-Dist: pytest~=8.3
 

nvidia_nat_test-1.3.0rc2.dist-info/RECORD

@@ -0,0 +1,16 @@
+nat/meta/pypi.md,sha256=LLKJHg5oN1-M9Pqfk3Bmphkk4O2TFsyiixuK5T0Y-gw,1100
+nat/test/__init__.py,sha256=_RnTJnsUucHvla_nYKqD4O4g8Bz0tcuDRzWk1bEhcy0,875
+nat/test/embedder.py,sha256=ClDyK1kna4hCBSlz71gK1B-ZjlwcBHTDQRekoNM81Bs,1809
+nat/test/functions.py,sha256=ZxXVzfaLBGOpR5qtmMrKU7q-M9-vVGGj3Xi5mrw4vHY,3557
+nat/test/llm.py,sha256=osJWGsJN7x-JGOaitueKeSwuJPVmnIFqJUCz28ngSRg,8215
+nat/test/memory.py,sha256=xki_A2yiMhEZuQk60K7t04QRqf32nQqnfzD5Iv7fkvw,1456
+nat/test/object_store_tests.py,sha256=PyJioOtoSzILPq6LuD-sOZ_89PIcgXWZweoHBQpK2zQ,4281
+nat/test/plugin.py,sha256=P6awHPN-iVmb8VIzLX_pd9rMnR4PD8Of3H_ypPPFr8Q,11246
+nat/test/register.py,sha256=o1BEA5fyxyFyCxXhQ6ArmtuNpgRyTEfvw6HdBgECPLI,897
+nat/test/tool_test_runner.py,sha256=SxavwXHkvCQDl_PUiiiqgvGfexKJJTeBdI5i1qk6AzI,21712
+nat/test/utils.py,sha256=y77p5uVpRPSQqVOnetBLvJVsSRgS4_fEgcuRoHwvRKE,3187
+nvidia_nat_test-1.3.0rc2.dist-info/METADATA,sha256=vyclpVYsAUUgANU3_ZbjWPqzGolQdTNjnVwq0_K-t5E,1608
+nvidia_nat_test-1.3.0rc2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+nvidia_nat_test-1.3.0rc2.dist-info/entry_points.txt,sha256=7dOP9XB6iMDqvav3gYx9VWUwA8RrFzhbAa8nGeC8e4Y,99
+nvidia_nat_test-1.3.0rc2.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
+nvidia_nat_test-1.3.0rc2.dist-info/RECORD,,

nvidia_nat_test-1.3.0.dev2.dist-info/RECORD

@@ -1,14 +0,0 @@
-nat/meta/pypi.md,sha256=LLKJHg5oN1-M9Pqfk3Bmphkk4O2TFsyiixuK5T0Y-gw,1100
-nat/test/__init__.py,sha256=_RnTJnsUucHvla_nYKqD4O4g8Bz0tcuDRzWk1bEhcy0,875
-nat/test/embedder.py,sha256=ClDyK1kna4hCBSlz71gK1B-ZjlwcBHTDQRekoNM81Bs,1809
-nat/test/functions.py,sha256=0ScrdsjcxCsPRLnyb5gfwukmvZxFi_ptCswLSIG0DVY,3095
-nat/test/memory.py,sha256=xki_A2yiMhEZuQk60K7t04QRqf32nQqnfzD5Iv7fkvw,1456
-nat/test/object_store_tests.py,sha256=PyJioOtoSzILPq6LuD-sOZ_89PIcgXWZweoHBQpK2zQ,4281
-nat/test/plugin.py,sha256=fp39ib0W63vfqX6Ssvq4sCuSd8Lm6yQyknL3_qRijgI,3610
-nat/test/register.py,sha256=jU1pW5wf20ZmCOTgkaQshKZfvYh8_-sMJ4P3xXilfTY,891
-nat/test/tool_test_runner.py,sha256=ccErldob2VwBbVL0_pmLrOcKLc18qYjxeAEACYoKKGQ,17469
-nvidia_nat_test-1.3.0.dev2.dist-info/METADATA,sha256=LFC2hgfp8ITqrKuES1NOCXbkyBwWPm_8OhjiIPIjxe4,1458
-nvidia_nat_test-1.3.0.dev2.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-nvidia_nat_test-1.3.0.dev2.dist-info/entry_points.txt,sha256=7dOP9XB6iMDqvav3gYx9VWUwA8RrFzhbAa8nGeC8e4Y,99
-nvidia_nat_test-1.3.0.dev2.dist-info/top_level.txt,sha256=8-CJ2cP6-f0ZReXe5Hzqp-5pvzzHz-5Ds5H2bGqh1-U,4
-nvidia_nat_test-1.3.0.dev2.dist-info/RECORD,,