From 180eddc800f59233fc232b7475b13f90f14f70f9 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 | 10 ++++- packages/toolbox-core/tests/test_tool.py | 39 ++++++++++++++++++- 2 files changed, 46 insertions(+), 3 deletions(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index 5351293a..cdfb7ef3 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -283,7 +283,7 @@ def add_auth_token_getters( new_getters = types.MappingProxyType( dict(self.__auth_service_token_getters, **auth_token_getters) ) - # create a read-only updated for params that are still required + # find the updated required authn params and the auth token getters used new_req_authn_params, new_req_authz_tokens, used_auth_token_getters = ( identify_required_authn_params( self.__required_authn_params, @@ -292,10 +292,16 @@ 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( auth_service_token_getters=new_getters, + # create a read-only map for params that are still required required_authn_params=types.MappingProxyType(new_req_authn_params), required_authz_tokens=new_req_authz_tokens, ) 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 913d9904203cf25a8798cbaedc8b2697c4eba1b6 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 | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/packages/toolbox-core/src/toolbox_core/tool.py b/packages/toolbox-core/src/toolbox_core/tool.py index cdfb7ef3..321352bc 100644 --- a/packages/toolbox-core/src/toolbox_core/tool.py +++ b/packages/toolbox-core/src/toolbox_core/tool.py @@ -283,7 +283,9 @@ def add_auth_token_getters( new_getters = types.MappingProxyType( dict(self.__auth_service_token_getters, **auth_token_getters) ) - # find the updated required authn params and the auth token getters used + + # 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_required_authn_params( self.__required_authn_params, From b7b0dc19c3dbff83857901886a847a8667acd956 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 54b925ee..2e4c14b6 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 4405195d47080f30260497db7545c7321051df34 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 2e4c14b6..19b13e98 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): """