Fix tests

Author: strawgate

pydantic_ai_slim/pydantic_ai/agent.py

@@ -1644,12 +1644,10 @@ def toolset(
 
         The decorator can be used bare (`agent.toolset`).
 
-        Overloads for every possible signature of `toolset` are included so the decorator doesn't obscure
-        the type of the function.
-
         Example:
         ```python
         from pydantic_ai import Agent, RunContext
+        from pydantic_ai.toolsets import AbstractToolset, FunctionToolset
 
         agent = Agent('test', deps_type=str)
 
@@ -1792,7 +1790,9 @@ def _prepare_output_schema(
 
         return schema  # pyright: ignore[reportReturnType]
 
-    async def _materialize_toolset_functions(self, run_context: RunContext[AgentDepsT]) -> list[AbstractToolset[AgentDepsT]]:
+    async def _materialize_toolset_functions(
+        self, run_context: RunContext[AgentDepsT]
+    ) -> list[AbstractToolset[AgentDepsT]]:
         materialized_toolsets: list[AbstractToolset[AgentDepsT]] = []
 
         for toolset_function in self._toolset_functions:
@@ -2356,7 +2356,6 @@ def _traceparent(self, *, required: bool = True) -> str | None:
     def data(self) -> OutputDataT:
         return self.output
 
-
     def _set_output_tool_return(self, return_content: str) -> list[_messages.ModelMessage]:
         """Set return content for the output tool.
 

pydantic_ai_slim/pydantic_ai/toolsets/abstract.py

@@ -1,8 +1,9 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Awaitable
 from dataclasses import dataclass
-from typing import TYPE_CHECKING, Any, Awaitable, Callable, Generic, Literal, Protocol
+from typing import TYPE_CHECKING, Any, Callable, Generic, Literal, Protocol
 
 from pydantic_core import SchemaValidator
 from typing_extensions import Self, TypeAlias
@@ -16,7 +17,9 @@
     from .prepared import PreparedToolset
     from .renamed import RenamedToolset
 
-ToolsetFunc: TypeAlias = 'Callable[[RunContext[AgentDepsT]], AbstractToolset[AgentDepsT] | Awaitable[AbstractToolset[AgentDepsT]]]'
+ToolsetFunc: TypeAlias = (
+    'Callable[[RunContext[AgentDepsT]], AbstractToolset[AgentDepsT] | Awaitable[AbstractToolset[AgentDepsT]]]'
+)
 """An sync/async function which takes a run context and returns a toolset."""
 
 

tests/test_agent.py

@@ -5,6 +5,8 @@
 from datetime import timezone
 from typing import Any, Callable, Union
 
+from pydantic_ai._agent_graph import build_run_context
+
 import httpx
 import pytest
 from dirty_equals import IsJson
@@ -3739,24 +3741,43 @@ def respond(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse:
     assert len(available_tools) == 1
     assert toolset_creation_count == 1
 
-def test_toolset_decorator():
+
+async def test_toolset_decorator():
     toolset = FunctionToolset()
 
     @toolset.tool
     def foo() -> str:
         return 'Hello from foo'
 
-
     agent = Agent('test')
 
     @agent.toolset
     def create_function_toolset(ctx: RunContext[None]) -> AbstractToolset[None]:
         return toolset
 
+    def create_function_toolset_bare(ctx: RunContext[None]) -> AbstractToolset[None]:
+        return toolset
+
+    agent.toolset(create_function_toolset_bare)
+
     agent_toolset_functions = agent._toolset_functions  # pyright: ignore[reportPrivateUsage]
 
-    assert len(agent_toolset_functions) == 1
+    assert len(agent_toolset_functions) == 2
     assert agent_toolset_functions[0] is create_function_toolset
+    assert agent_toolset_functions[1] is create_function_toolset_bare
+
+    fake_run_context = RunContext(
+        deps=None,
+        model=TestModel(),
+        usage=Usage(),
+        prompt=None,
+        messages=[],
+        run_step=0,
+    )
+
+    toolsets = await agent._materialize_toolset_functions(run_context=fake_run_context)  # pyright: ignore[reportPrivateUsage]
+    assert len(toolsets) == 2
+
 
 def test_adding_tools_during_run():
     toolset = FunctionToolset()