fix

Author: loveyana

tests/test_ve_identity_function_tool.py

@@ -27,9 +27,10 @@
 class TestVeIdentityFunctionToolInit:
     """Tests for VeIdentityFunctionTool initialization."""
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_api_key_auth(self, mock_identity_client):
         """Test initializing with API key auth config."""
+
         async def test_func(api_key: str):
             return f"Called with {api_key}"
 
@@ -41,26 +42,24 @@ async def test_func(api_key: str):
         assert tool._auth_config == config
         assert tool._into == "api_key"
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_oauth2_auth(self, mock_identity_client):
         """Test initializing with OAuth2 auth config."""
+
         async def test_func(access_token: str):
             return f"Called with {access_token}"
 
-        config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="M2M"
-        )
+        config = oauth2_auth(provider_name="github", scopes=["repo"], auth_flow="M2M")
         tool = VeIdentityFunctionTool(func=test_func, auth_config=config)
 
         assert tool.func == test_func
         assert tool._auth_config == config
         assert tool._into == "access_token"
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_workload_auth(self, mock_identity_client):
         """Test initializing with workload auth config."""
+
         async def test_func(access_token: str):
             return f"Called with {access_token}"
 
@@ -71,23 +70,23 @@ async def test_func(access_token: str):
         assert tool._auth_config == config
         assert tool._into == "access_token"
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_custom_into_parameter(self, mock_identity_client):
         """Test initializing with custom 'into' parameter."""
+
         async def test_func(custom_token: str):
             return f"Called with {custom_token}"
 
         config = api_key_auth("test-provider")
         tool = VeIdentityFunctionTool(
-            func=test_func,
-            auth_config=config,
-            into="custom_token"
+            func=test_func, auth_config=config, into="custom_token"
         )
 
         assert tool._into == "custom_token"
 
     def test_init_with_unsupported_auth_config(self):
         """Test that unsupported auth config raises ValueError."""
+
         async def test_func(token: str):
             return f"Called with {token}"
 
@@ -103,9 +102,10 @@ class TestVeIdentityFunctionToolRunAsync:
     """Tests for VeIdentityFunctionTool.run_async method."""
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     async def test_run_async_with_api_key(self, mock_identity_client):
         """Test run_async with API key authentication."""
+
         async def test_func(api_key: str):
             return f"Result: {api_key}"
 
@@ -122,17 +122,14 @@ async def test_func(api_key: str):
         tool.run_with_identity_auth.assert_called_once()
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     async def test_run_async_with_oauth2(self, mock_identity_client):
         """Test run_async with OAuth2 authentication."""
+
         async def test_func(access_token: str):
             return f"Result: {access_token}"
 
-        config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="M2M"
-        )
+        config = oauth2_auth(provider_name="github", scopes=["repo"], auth_flow="M2M")
         tool = VeIdentityFunctionTool(func=test_func, auth_config=config)
 
         # Mock the run_with_identity_auth method
@@ -145,18 +142,18 @@ async def test_func(access_token: str):
         tool.run_with_identity_auth.assert_called_once()
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
-    async def test_run_async_handles_auth_required_exception(self, mock_identity_client):
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
+    async def test_run_async_handles_auth_required_exception(
+        self, mock_identity_client
+    ):
         """Test that run_async handles AuthRequiredException."""
         from veadk.integrations.ve_identity.auth_mixins import AuthRequiredException
 
         async def test_func(access_token: str):
             return f"Result: {access_token}"
 
         config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="USER_FEDERATION"
+            provider_name="github", scopes=["repo"], auth_flow="USER_FEDERATION"
         )
         tool = VeIdentityFunctionTool(func=test_func, auth_config=config)
 
@@ -174,9 +171,10 @@ class TestVeIdentityFunctionToolExecuteWithCredential:
     """Tests for VeIdentityFunctionTool._execute_with_credential method."""
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     async def test_execute_with_credential_injects_api_key(self, mock_identity_client):
         """Test that _execute_with_credential injects API key."""
+
         async def test_func(api_key: str):
             return f"Result: {api_key}"
 
@@ -188,32 +186,31 @@ async def test_func(api_key: str):
         credential.api_key = "test-api-key"
 
         # Mock parent's run_async
-        with patch.object(tool.__class__.__bases__[1], 'run_async', new_callable=AsyncMock) as mock_run:
+        with patch.object(
+            tool.__class__.__bases__[1], "run_async", new_callable=AsyncMock
+        ) as mock_run:
             mock_run.return_value = "Result: test-api-key"
 
             tool_context = Mock()
-            result = await tool._execute_with_credential(
-                args={},
-                tool_context=tool_context,
-                credential=credential
+            await tool._execute_with_credential(
+                args={}, tool_context=tool_context, credential=credential
             )
 
             # Verify that run_async was called with injected api_key
             call_args = mock_run.call_args
-            assert call_args[1]['args']['api_key'] == "test-api-key"
+            assert call_args[1]["args"]["api_key"] == "test-api-key"
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
-    async def test_execute_with_credential_injects_oauth2_token(self, mock_identity_client):
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
+    async def test_execute_with_credential_injects_oauth2_token(
+        self, mock_identity_client
+    ):
         """Test that _execute_with_credential injects OAuth2 access token."""
+
         async def test_func(access_token: str):
             return f"Result: {access_token}"
 
-        config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="M2M"
-        )
+        config = oauth2_auth(provider_name="github", scopes=["repo"], auth_flow="M2M")
         tool = VeIdentityFunctionTool(func=test_func, auth_config=config)
 
         # Mock credential
@@ -222,17 +219,16 @@ async def test_func(access_token: str):
         credential.oauth2.access_token = "test-oauth2-token"
 
         # Mock parent's run_async
-        with patch.object(tool.__class__.__bases__[1], 'run_async', new_callable=AsyncMock) as mock_run:
+        with patch.object(
+            tool.__class__.__bases__[1], "run_async", new_callable=AsyncMock
+        ) as mock_run:
             mock_run.return_value = "Result: test-oauth2-token"
 
             tool_context = Mock()
-            result = await tool._execute_with_credential(
-                args={},
-                tool_context=tool_context,
-                credential=credential
+            await tool._execute_with_credential(
+                args={}, tool_context=tool_context, credential=credential
             )
 
             # Verify that run_async was called with injected access_token
             call_args = mock_run.call_args
-            assert call_args[1]['args']['access_token'] == "test-oauth2-token"
-
+            assert call_args[1]["args"]["access_token"] == "test-oauth2-token"

tests/test_ve_identity_mcp_tool.py

@@ -1,3 +1,18 @@
+# Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd. and/or its affiliates.
+#
+# 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.
+
+
 """Unit tests for VeIdentityMcpTool."""
 
 import pytest
@@ -12,60 +27,56 @@
 class TestVeIdentityMcpToolInit:
     """Tests for VeIdentityMcpTool initialization."""
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_api_key_auth(self, mock_identity_client):
         """Test initializing with API key auth config."""
         mcp_tool = Mock()
         mcp_tool.name = "test_tool"
         mcp_tool.description = "Test tool description"
-        
+
         mcp_session_manager = Mock()
         config = api_key_auth("test-provider")
-        
+
         tool = VeIdentityMcpTool(
             mcp_tool=mcp_tool,
             mcp_session_manager=mcp_session_manager,
-            auth_config=config
+            auth_config=config,
         )
-        
+
         assert tool.name == "test_tool"
         assert tool.description == "Test tool description"
         assert tool._mcp_tool == mcp_tool
         assert tool._mcp_session_manager == mcp_session_manager
 
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     def test_init_with_oauth2_auth(self, mock_identity_client):
         """Test initializing with OAuth2 auth config."""
         mcp_tool = Mock()
         mcp_tool.name = "github_tool"
         mcp_tool.description = "GitHub tool"
-        
+
         mcp_session_manager = Mock()
-        config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="M2M"
-        )
-        
+        config = oauth2_auth(provider_name="github", scopes=["repo"], auth_flow="M2M")
+
         tool = VeIdentityMcpTool(
             mcp_tool=mcp_tool,
             mcp_session_manager=mcp_session_manager,
-            auth_config=config
+            auth_config=config,
         )
-        
+
         assert tool.name == "github_tool"
         assert tool._auth_config == config
 
     def test_init_with_none_mcp_tool(self):
         """Test that initialization fails with None mcp_tool."""
         mcp_session_manager = Mock()
         config = api_key_auth("test-provider")
-        
+
         with pytest.raises(ValueError, match="mcp_tool cannot be None"):
             VeIdentityMcpTool(
                 mcp_tool=None,
                 mcp_session_manager=mcp_session_manager,
-                auth_config=config
+                auth_config=config,
             )
 
     def test_init_with_none_mcp_session_manager(self):
@@ -74,71 +85,68 @@ def test_init_with_none_mcp_session_manager(self):
         mcp_tool.name = "test_tool"
         mcp_tool.description = "Test"
         config = api_key_auth("test-provider")
-        
+
         with pytest.raises(ValueError, match="mcp_session_manager cannot be None"):
             VeIdentityMcpTool(
-                mcp_tool=mcp_tool,
-                mcp_session_manager=None,
-                auth_config=config
+                mcp_tool=mcp_tool, mcp_session_manager=None, auth_config=config
             )
 
 
 class TestVeIdentityMcpToolRunAsync:
     """Tests for VeIdentityMcpTool.run_async method."""
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
     async def test_run_async_with_api_key(self, mock_identity_client):
         """Test run_async with API key authentication."""
         mcp_tool = Mock()
         mcp_tool.name = "test_tool"
         mcp_tool.description = "Test"
-        
+
         mcp_session_manager = Mock()
         config = api_key_auth("test-provider")
-        
+
         tool = VeIdentityMcpTool(
             mcp_tool=mcp_tool,
             mcp_session_manager=mcp_session_manager,
-            auth_config=config
+            auth_config=config,
         )
-        
+
         # Mock the run_with_identity_auth method
         tool.run_with_identity_auth = AsyncMock(return_value="Result: test-key")
-        
+
         tool_context = Mock()
         result = await tool.run_async(args={}, tool_context=tool_context)
-        
+
         assert result == "Result: test-key"
         tool.run_with_identity_auth.assert_called_once()
 
     @pytest.mark.asyncio
-    @patch('veadk.integrations.ve_identity.auth_mixins.IdentityClient')
-    async def test_run_async_handles_auth_required_exception(self, mock_identity_client):
+    @patch("veadk.integrations.ve_identity.auth_mixins.IdentityClient")
+    async def test_run_async_handles_auth_required_exception(
+        self, mock_identity_client
+    ):
         """Test that run_async handles AuthRequiredException."""
         mcp_tool = Mock()
         mcp_tool.name = "test_tool"
         mcp_tool.description = "Test"
-        
+
         mcp_session_manager = Mock()
         config = oauth2_auth(
-            provider_name="github",
-            scopes=["repo"],
-            auth_flow="USER_FEDERATION"
+            provider_name="github", scopes=["repo"], auth_flow="USER_FEDERATION"
         )
-        
+
         tool = VeIdentityMcpTool(
             mcp_tool=mcp_tool,
             mcp_session_manager=mcp_session_manager,
-            auth_config=config
+            auth_config=config,
         )
-        
+
         # Mock the run_with_identity_auth to raise AuthRequiredException
         auth_exception = AuthRequiredException("Please authorize")
         tool.run_with_identity_auth = AsyncMock(side_effect=auth_exception)
-        
+
         tool_context = Mock()
         result = await tool.run_async(args={}, tool_context=tool_context)
-        
-        assert result == "Please authorize"
 
+        assert result == "Please authorize"