From f3c384f2376eb12cc4d7b60d9e2797c9fb0b460f Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Mon, 5 May 2025 20:45:56 +0530 Subject: [PATCH 1/4] fix: Add validation to ensure added auth token getters are used by the tool Previously, the `ToolboxTool.add_auth_token_getters` method only validated against existing registered getters or conflicts with client headers. It did not verify if *all* the auth token getters provided were actually used or required by the specific tool instance they were being added to. This PR enhances the validation in `add_auth_token_getters`. It now leverages the `used_auth_token_getters` information returned by the existing `identify_required_authn_params` call. This allows the method to confirm that every getter passed in is genuinely required by the tool, raising an error if any are unused. This ensures that only relevant auth token getters are attempted to be registered for a tool, preventing misconfigurations and human errors. > [!NOTE] > This validation aligns with the existing validation logic already present in the `ToolboxClient.load_tool` method, promoting a consistent approach to handling auth token getter requirements across the codebase. --- .../toolbox-core/src/toolbox_core/tool.py | 9 ++++- packages/toolbox-core/tests/test_tool.py | 39 ++++++++++++++++++- 2 files changed, 45 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index a0e5eb2c..14f17384 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -309,7 +309,7 @@ def add_auth_token_getters( new_getters = dict(self.__auth_service_token_getters, **auth_token_getters) - # find the updated requirements + # find the updated auth requirements new_req_authn_params, new_req_authz_tokens, used_auth_token_getters = ( identify_auth_requirements( self.__required_authn_params, @@ -318,7 +318,12 @@ def add_auth_token_getters( ) ) - # TODO: Add validation for used_auth_token_getters + # ensure no auth token getter provided remains unused + unused_auth = set(incoming_services) - used_auth_token_getters + if unused_auth: + raise ValueError( + f"Authentication source(s) `{', '.join(unused_auth)}` unused by tool `{self.__name__}`." + ) return self.__copy( # create a read-only map for updated getters, params and tokens that are still required diff --git a/packages/toolbox-core/tests/test_tool.py b/packages/toolbox-core/tests/test_tool.py index cbf64efd..4b723883 100644 --- a/packages/toolbox-core/tests/test_tool.py +++ b/packages/toolbox-core/tests/test_tool.py @@ -11,8 +11,10 @@ # 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 inspect -from typing import AsyncGenerator, Callable +from typing import AsyncGenerator, Callable, Mapping from unittest.mock import AsyncMock, Mock import pytest @@ -92,6 +94,12 @@ def auth_header_key() -> str: return "test-auth_token" +@pytest.fixture +def unused_auth_getters() -> dict[str, Callable[[], str]]: + """Provides an auth getter for a service not required by sample_tool.""" + return {"unused-auth-service": lambda: "unused-token-value"} + + def test_create_func_docstring_one_param_real_schema(): """ Tests create_func_docstring with one real ParameterSchema instance. @@ -432,3 +440,32 @@ def test_tool_add_auth_token_getters_conflict_with_existing_client_header( with pytest.raises(ValueError, match=expected_error_message): tool_instance.add_auth_token_getters(new_auth_getters_causing_conflict) + + +def test_add_auth_token_getters_unused_token( + http_session: ClientSession, + sample_tool_params: list[ParameterSchema], + sample_tool_description: str, + unused_auth_getters: Mapping[str, Callable[[], str]], +): + """ + Tests ValueError when add_auth_token_getters is called with a getter for + an unused authentication service. + """ + tool_instance = ToolboxTool( + session=http_session, + base_url=TEST_BASE_URL, + name=TEST_TOOL_NAME, + description=sample_tool_description, + params=sample_tool_params, + required_authn_params={}, + required_authz_tokens=[], + auth_service_token_getters={}, + bound_params={}, + client_headers={}, + ) + + expected_error_message = "Authentication source\(s\) \`unused-auth-service\` unused by tool \`sample_tool\`." + + with pytest.raises(ValueError, match=expected_error_message): + tool_instance.add_auth_token_getters(unused_auth_getters) From f6d166102cbae481174e22ec5ec0f4d986f0d141 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Tue, 6 May 2025 16:59:38 +0530 Subject: [PATCH 2/4] docs: Improve helper comment --- packages/toolbox-core/src/toolbox_core/tool.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index 14f17384..2bbb1a3f 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -309,7 +309,8 @@ def add_auth_token_getters( new_getters = dict(self.__auth_service_token_getters, **auth_token_getters) - # find the updated auth requirements + # find the updated required authn params, authz tokens and the auth + # token getters used new_req_authn_params, new_req_authz_tokens, used_auth_token_getters = ( identify_auth_requirements( self.__required_authn_params, From 91f998899f072f7c8ccd876dccc5b693e2535fea Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Tue, 6 May 2025 17:11:47 +0530 Subject: [PATCH 3/4] chore: Add unit test cases --- packages/toolbox-core/tests/test_client.py | 29 ++++++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index 57f8d56c..89e517c3 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -372,6 +372,35 @@ async def test_add_auth_token_getters_duplicate_fail(self, tool_name, client): authed_tool.add_auth_token_getters({AUTH_SERVICE: {}}) + @pytest.mark.asyncio + async def test_add_auth_token_getters_missing_fail(self, tool_name, client): + """ + Tests that adding a missing auth token getter raises ValueError. + """ + AUTH_SERVICE = "xmy-auth-service" + + tool = await client.load_tool(tool_name) + + with pytest.raises( + ValueError, + match=f"Authentication source\(s\) \`{AUTH_SERVICE}\` unused by tool \`{tool_name}\`.", + ): + tool.add_auth_token_getters({AUTH_SERVICE: {}}) + + @pytest.mark.asyncio + async def test_constructor_getters_missing_fail(self, tool_name, client): + """ + Tests that adding a missing auth token getter raises ValueError. + """ + AUTH_SERVICE = "xmy-auth-service" + + with pytest.raises( + ValueError, + match=f"Validation failed for tool '{tool_name}': unused auth tokens: {AUTH_SERVICE}.", + ): + await client.load_tool(tool_name, auth_token_getters={AUTH_SERVICE: {}}) + + class TestBoundParameter: @pytest.fixture From 88175d7e15c51b7c05f966d769053a4b44b9d5b5 Mon Sep 17 00:00:00 2001 From: Anubhav Dhawan Date: Tue, 6 May 2025 17:12:37 +0530 Subject: [PATCH 4/4] chore: Delint --- packages/toolbox-core/tests/test_client.py | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/toolbox-core/tests/test_client.py b/packages/toolbox-core/tests/test_client.py index 89e517c3..1d7aedbd 100644 --- a/packages/toolbox-core/tests/test_client.py +++ b/packages/toolbox-core/tests/test_client.py @@ -371,7 +371,6 @@ async def test_add_auth_token_getters_duplicate_fail(self, tool_name, client): ): authed_tool.add_auth_token_getters({AUTH_SERVICE: {}}) - @pytest.mark.asyncio async def test_add_auth_token_getters_missing_fail(self, tool_name, client): """