From 3592240d142d9b3a2954657d84d968b6f61327a3 Mon Sep 17 00:00:00 2001 From: Novice Date: Wed, 17 Sep 2025 10:29:22 +0800 Subject: [PATCH] chore: fix db container test --- .../tools/test_mcp_tools_manage_service.py | 378 ++++++++++++------ 1 file changed, 264 insertions(+), 114 deletions(-) diff --git a/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py b/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py index dd22dcbfd1..cae8906736 100644 --- a/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py +++ b/api/tests/test_containers_integration_tests/services/tools/test_mcp_tools_manage_service.py @@ -20,12 +20,21 @@ class TestMCPToolManageService: patch("services.tools.mcp_tools_manage_service.ToolTransformService") as mock_tool_transform_service, ): # Setup default mock returns + from core.tools.entities.api_entities import ToolProviderApiEntity + from core.tools.entities.common_entities import I18nObject + mock_encrypter.encrypt_token.return_value = "encrypted_server_url" - mock_tool_transform_service.mcp_provider_to_user_provider.return_value = { - "id": "test_id", - "name": "test_name", - "type": ToolProviderType.MCP, - } + mock_tool_transform_service.mcp_provider_to_user_provider.return_value = ToolProviderApiEntity( + id="test_id", + author="test_author", + name="test_name", + type=ToolProviderType.MCP, + description=I18nObject(en_US="Test Description", zh_Hans="测试描述"), + icon={"type": "emoji", "content": "🤖"}, + label=I18nObject(en_US="Test Label", zh_Hans="测试标签"), + labels=[], + tools=[], + ) yield { "encrypter": mock_encrypter, @@ -144,7 +153,10 @@ class TestMCPToolManageService: ) # Act: Execute the method under test - result = MCPToolManageService.get_mcp_provider_by_provider_id(mcp_provider.id, tenant.id) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.get_provider(provider_id=mcp_provider.id, tenant_id=tenant.id) # Assert: Verify the expected outcomes assert result is not None @@ -154,8 +166,6 @@ class TestMCPToolManageService: assert result.user_id == account.id # Verify database state - from extensions.ext_database import db - db.session.refresh(result) assert result.id is not None assert result.server_identifier == mcp_provider.server_identifier @@ -177,11 +187,14 @@ class TestMCPToolManageService: db_session_with_containers, mock_external_service_dependencies ) - non_existent_id = fake.uuid4() + non_existent_id = str(fake.uuid4()) # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.get_mcp_provider_by_provider_id(non_existent_id, tenant.id) + service.get_provider(provider_id=non_existent_id, tenant_id=tenant.id) def test_get_mcp_provider_by_provider_id_tenant_isolation( self, db_session_with_containers, mock_external_service_dependencies @@ -210,8 +223,11 @@ class TestMCPToolManageService: ) # Act & Assert: Verify tenant isolation + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.get_mcp_provider_by_provider_id(mcp_provider1.id, tenant2.id) + service.get_provider(provider_id=mcp_provider1.id, tenant_id=tenant2.id) def test_get_mcp_provider_by_server_identifier_success( self, db_session_with_containers, mock_external_service_dependencies @@ -235,7 +251,10 @@ class TestMCPToolManageService: ) # Act: Execute the method under test - result = MCPToolManageService.get_mcp_provider_by_server_identifier(mcp_provider.server_identifier, tenant.id) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.get_provider(server_identifier=mcp_provider.server_identifier, tenant_id=tenant.id) # Assert: Verify the expected outcomes assert result is not None @@ -245,8 +264,6 @@ class TestMCPToolManageService: assert result.user_id == account.id # Verify database state - from extensions.ext_database import db - db.session.refresh(result) assert result.id is not None assert result.name == mcp_provider.name @@ -268,11 +285,14 @@ class TestMCPToolManageService: db_session_with_containers, mock_external_service_dependencies ) - non_existent_identifier = fake.uuid4() + non_existent_identifier = str(fake.uuid4()) # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.get_mcp_provider_by_server_identifier(non_existent_identifier, tenant.id) + service.get_provider(server_identifier=non_existent_identifier, tenant_id=tenant.id) def test_get_mcp_provider_by_server_identifier_tenant_isolation( self, db_session_with_containers, mock_external_service_dependencies @@ -301,8 +321,11 @@ class TestMCPToolManageService: ) # Act & Assert: Verify tenant isolation + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.get_mcp_provider_by_server_identifier(mcp_provider1.server_identifier, tenant2.id) + service.get_provider(server_identifier=mcp_provider1.server_identifier, tenant_id=tenant2.id) def test_create_mcp_provider_success(self, db_session_with_containers, mock_external_service_dependencies): """ @@ -322,15 +345,29 @@ class TestMCPToolManageService: ) # Setup mocks for provider creation + from core.tools.entities.api_entities import ToolProviderApiEntity + from core.tools.entities.common_entities import I18nObject + mock_external_service_dependencies["encrypter"].encrypt_token.return_value = "encrypted_server_url" - mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.return_value = { - "id": "new_provider_id", - "name": "Test MCP Provider", - "type": ToolProviderType.MCP, - } + mock_external_service_dependencies[ + "tool_transform_service" + ].mcp_provider_to_user_provider.return_value = ToolProviderApiEntity( + id="new_provider_id", + author=account.name, + name="Test MCP Provider", + type=ToolProviderType.MCP, + description=I18nObject(en_US="Test MCP Provider Description", zh_Hans="测试MCP提供者描述"), + icon={"type": "emoji", "content": "🤖"}, + label=I18nObject(en_US="Test MCP Provider", zh_Hans="测试MCP提供者"), + labels=[], + tools=[], + ) # Act: Execute the method under test - result = MCPToolManageService.create_mcp_provider( + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.create_provider( tenant_id=tenant.id, name="Test MCP Provider", server_url="https://example.com/mcp", @@ -345,8 +382,8 @@ class TestMCPToolManageService: # Assert: Verify the expected outcomes assert result is not None - assert result["name"] == "Test MCP Provider" - assert result["type"] == ToolProviderType.MCP + assert result.name == "Test MCP Provider" + assert result.type == ToolProviderType.MCP # Verify database state from extensions.ext_database import db @@ -386,7 +423,10 @@ class TestMCPToolManageService: ) # Create first provider - MCPToolManageService.create_mcp_provider( + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider", server_url="https://example1.com/mcp", @@ -401,7 +441,7 @@ class TestMCPToolManageService: # Act & Assert: Verify proper error handling for duplicate name with pytest.raises(ValueError, match="MCP tool Test MCP Provider already exists"): - MCPToolManageService.create_mcp_provider( + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider", # Duplicate name server_url="https://example2.com/mcp", @@ -432,7 +472,10 @@ class TestMCPToolManageService: ) # Create first provider - MCPToolManageService.create_mcp_provider( + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider 1", server_url="https://example.com/mcp", @@ -446,8 +489,8 @@ class TestMCPToolManageService: ) # Act & Assert: Verify proper error handling for duplicate server URL - with pytest.raises(ValueError, match="MCP tool https://example.com/mcp already exists"): - MCPToolManageService.create_mcp_provider( + with pytest.raises(ValueError, match="MCP tool with this server URL already exists"): + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider 2", server_url="https://example.com/mcp", # Duplicate URL @@ -478,7 +521,10 @@ class TestMCPToolManageService: ) # Create first provider - MCPToolManageService.create_mcp_provider( + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider 1", server_url="https://example1.com/mcp", @@ -493,7 +539,7 @@ class TestMCPToolManageService: # Act & Assert: Verify proper error handling for duplicate server identifier with pytest.raises(ValueError, match="MCP tool test_identifier_123 already exists"): - MCPToolManageService.create_mcp_provider( + service.create_provider( tenant_id=tenant.id, name="Test MCP Provider 2", server_url="https://example2.com/mcp", @@ -543,23 +589,59 @@ class TestMCPToolManageService: db.session.commit() # Setup mock for transformation service + from core.tools.entities.api_entities import ToolProviderApiEntity + from core.tools.entities.common_entities import I18nObject + mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.side_effect = [ - {"id": provider1.id, "name": provider1.name, "type": ToolProviderType.MCP}, - {"id": provider2.id, "name": provider2.name, "type": ToolProviderType.MCP}, - {"id": provider3.id, "name": provider3.name, "type": ToolProviderType.MCP}, + ToolProviderApiEntity( + id=provider1.id, + author=account.name, + name=provider1.name, + type=ToolProviderType.MCP, + description=I18nObject(en_US="Alpha Provider Description", zh_Hans="Alpha提供者描述"), + icon={"type": "emoji", "content": "🅰️"}, + label=I18nObject(en_US=provider1.name, zh_Hans=provider1.name), + labels=[], + tools=[], + ), + ToolProviderApiEntity( + id=provider2.id, + author=account.name, + name=provider2.name, + type=ToolProviderType.MCP, + description=I18nObject(en_US="Beta Provider Description", zh_Hans="Beta提供者描述"), + icon={"type": "emoji", "content": "🅱️"}, + label=I18nObject(en_US=provider2.name, zh_Hans=provider2.name), + labels=[], + tools=[], + ), + ToolProviderApiEntity( + id=provider3.id, + author=account.name, + name=provider3.name, + type=ToolProviderType.MCP, + description=I18nObject(en_US="Gamma Provider Description", zh_Hans="Gamma提供者描述"), + icon={"type": "emoji", "content": "Γ"}, + label=I18nObject(en_US=provider3.name, zh_Hans=provider3.name), + labels=[], + tools=[], + ), ] # Act: Execute the method under test - result = MCPToolManageService.retrieve_mcp_tools(tenant.id, for_list=True) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.list_providers(tenant_id=tenant.id, for_list=True) # Assert: Verify the expected outcomes assert result is not None assert len(result) == 3 # Verify correct ordering by name - assert result[0]["name"] == "Alpha Provider" - assert result[1]["name"] == "Beta Provider" - assert result[2]["name"] == "Gamma Provider" + assert result[0].name == "Alpha Provider" + assert result[1].name == "Beta Provider" + assert result[2].name == "Gamma Provider" # Verify mock interactions assert ( @@ -584,7 +666,10 @@ class TestMCPToolManageService: # No MCP providers created for this tenant # Act: Execute the method under test - result = MCPToolManageService.retrieve_mcp_tools(tenant.id, for_list=False) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.list_providers(tenant_id=tenant.id, for_list=False) # Assert: Verify the expected outcomes assert result is not None @@ -624,20 +709,46 @@ class TestMCPToolManageService: ) # Setup mock for transformation service + from core.tools.entities.api_entities import ToolProviderApiEntity + from core.tools.entities.common_entities import I18nObject + mock_external_service_dependencies["tool_transform_service"].mcp_provider_to_user_provider.side_effect = [ - {"id": provider1.id, "name": provider1.name, "type": ToolProviderType.MCP}, - {"id": provider2.id, "name": provider2.name, "type": ToolProviderType.MCP}, + ToolProviderApiEntity( + id=provider1.id, + author=account1.name, + name=provider1.name, + type=ToolProviderType.MCP, + description=I18nObject(en_US="Provider 1 Description", zh_Hans="提供者1描述"), + icon={"type": "emoji", "content": "1️⃣"}, + label=I18nObject(en_US=provider1.name, zh_Hans=provider1.name), + labels=[], + tools=[], + ), + ToolProviderApiEntity( + id=provider2.id, + author=account2.name, + name=provider2.name, + type=ToolProviderType.MCP, + description=I18nObject(en_US="Provider 2 Description", zh_Hans="提供者2描述"), + icon={"type": "emoji", "content": "2️⃣"}, + label=I18nObject(en_US=provider2.name, zh_Hans=provider2.name), + labels=[], + tools=[], + ), ] # Act: Execute the method under test for both tenants - result1 = MCPToolManageService.retrieve_mcp_tools(tenant1.id, for_list=True) - result2 = MCPToolManageService.retrieve_mcp_tools(tenant2.id, for_list=True) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result1 = service.list_providers(tenant_id=tenant1.id, for_list=True) + result2 = service.list_providers(tenant_id=tenant2.id, for_list=True) # Assert: Verify tenant isolation assert len(result1) == 1 assert len(result2) == 1 - assert result1[0]["id"] == provider1.id - assert result2[0]["id"] == provider2.id + assert result1[0].id == provider1.id + assert result2[0].id == provider2.id def test_list_mcp_tool_from_remote_server_success( self, db_session_with_containers, mock_external_service_dependencies @@ -661,17 +772,20 @@ class TestMCPToolManageService: mcp_provider = self._create_test_mcp_provider( db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id ) - mcp_provider.server_url = "encrypted_server_url" - mcp_provider.authed = False + # Use a valid base64 encoded string to avoid decryption errors + import base64 + + mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode() + mcp_provider.authed = True # Provider must be authenticated to list tools mcp_provider.tools = "[]" from extensions.ext_database import db db.session.commit() - # Mock the decrypted_server_url property to avoid encryption issues - with patch("models.tools.encrypter") as mock_encrypter: - mock_encrypter.decrypt_token.return_value = "https://example.com/mcp" + # Mock the decryption process at the rsa level to avoid key file issues + with patch("libs.rsa.decrypt") as mock_decrypt: + mock_decrypt.return_value = "https://example.com/mcp" # Mock MCPClient and its context manager mock_tools = [ @@ -683,13 +797,16 @@ class TestMCPToolManageService: )(), ] - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: # Setup mock client mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.return_value = mock_tools # Act: Execute the method under test - result = MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id) + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id) # Assert: Verify the expected outcomes assert result is not None @@ -705,16 +822,8 @@ class TestMCPToolManageService: assert mcp_provider.updated_at is not None # Verify mock interactions - mock_mcp_client.assert_called_once_with( - "https://example.com/mcp", - mcp_provider.id, - tenant.id, - authed=False, - for_list=True, - headers={}, - timeout=30.0, - sse_read_timeout=300.0, - ) + # MCPClientWithAuthRetry is called with different parameters + mock_mcp_client.assert_called_once() def test_list_mcp_tool_from_remote_server_auth_error( self, db_session_with_containers, mock_external_service_dependencies @@ -737,7 +846,10 @@ class TestMCPToolManageService: mcp_provider = self._create_test_mcp_provider( db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id ) - mcp_provider.server_url = "encrypted_server_url" + # Use a valid base64 encoded string to avoid decryption errors + import base64 + + mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode() mcp_provider.authed = False mcp_provider.tools = "[]" @@ -745,20 +857,23 @@ class TestMCPToolManageService: db.session.commit() - # Mock the decrypted_server_url property to avoid encryption issues - with patch("models.tools.encrypter") as mock_encrypter: - mock_encrypter.decrypt_token.return_value = "https://example.com/mcp" + # Mock the decryption process at the rsa level to avoid key file issues + with patch("libs.rsa.decrypt") as mock_decrypt: + mock_decrypt.return_value = "https://example.com/mcp" # Mock MCPClient to raise authentication error - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: from core.mcp.error import MCPAuthError mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.side_effect = MCPAuthError("Authentication required") # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="Please auth the tool first"): - MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id) + service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id) # Verify database state was not changed db.session.refresh(mcp_provider) @@ -786,32 +901,38 @@ class TestMCPToolManageService: mcp_provider = self._create_test_mcp_provider( db_session_with_containers, mock_external_service_dependencies, tenant.id, account.id ) - mcp_provider.server_url = "encrypted_server_url" - mcp_provider.authed = False + # Use a valid base64 encoded string to avoid decryption errors + import base64 + + mcp_provider.server_url = base64.b64encode(b"encrypted_server_url").decode() + mcp_provider.authed = True # Provider must be authenticated to test connection errors mcp_provider.tools = "[]" from extensions.ext_database import db db.session.commit() - # Mock the decrypted_server_url property to avoid encryption issues - with patch("models.tools.encrypter") as mock_encrypter: - mock_encrypter.decrypt_token.return_value = "https://example.com/mcp" + # Mock the decryption process at the rsa level to avoid key file issues + with patch("libs.rsa.decrypt") as mock_decrypt: + mock_decrypt.return_value = "https://example.com/mcp" # Mock MCPClient to raise connection error - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: from core.mcp.error import MCPError mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.side_effect = MCPError("Connection failed") # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="Failed to connect to MCP server: Connection failed"): - MCPToolManageService.list_mcp_tool_from_remote_server(tenant.id, mcp_provider.id) + service.list_provider_tools(tenant_id=tenant.id, provider_id=mcp_provider.id) # Verify database state was not changed db.session.refresh(mcp_provider) - assert mcp_provider.authed is False + assert mcp_provider.authed is True # Provider remains authenticated assert mcp_provider.tools == "[]" def test_delete_mcp_tool_success(self, db_session_with_containers, mock_external_service_dependencies): @@ -840,7 +961,8 @@ class TestMCPToolManageService: assert db.session.query(MCPToolProvider).filter_by(id=mcp_provider.id).first() is not None # Act: Execute the method under test - MCPToolManageService.delete_mcp_tool(tenant.id, mcp_provider.id) + service = MCPToolManageService(db.session()) + service.delete_provider(tenant_id=tenant.id, provider_id=mcp_provider.id) # Assert: Verify the expected outcomes # Provider should be deleted from database @@ -862,11 +984,14 @@ class TestMCPToolManageService: db_session_with_containers, mock_external_service_dependencies ) - non_existent_id = fake.uuid4() + non_existent_id = str(fake.uuid4()) # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.delete_mcp_tool(tenant.id, non_existent_id) + service.delete_provider(tenant_id=tenant.id, provider_id=non_existent_id) def test_delete_mcp_tool_tenant_isolation(self, db_session_with_containers, mock_external_service_dependencies): """ @@ -893,8 +1018,11 @@ class TestMCPToolManageService: ) # Act & Assert: Verify tenant isolation + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool not found"): - MCPToolManageService.delete_mcp_tool(tenant2.id, mcp_provider1.id) + service.delete_provider(tenant_id=tenant2.id, provider_id=mcp_provider1.id) # Verify provider still exists in tenant1 from extensions.ext_database import db @@ -929,7 +1057,8 @@ class TestMCPToolManageService: db.session.commit() # Act: Execute the method under test - MCPToolManageService.update_mcp_provider( + service = MCPToolManageService(db.session()) + service.update_provider( tenant_id=tenant.id, provider_id=mcp_provider.id, name="Updated MCP Provider", @@ -985,7 +1114,7 @@ class TestMCPToolManageService: db.session.commit() # Mock the reconnection method - with patch.object(MCPToolManageService, "_re_connect_mcp_provider") as mock_reconnect: + with patch.object(MCPToolManageService, "_reconnect_provider") as mock_reconnect: mock_reconnect.return_value = { "authed": True, "tools": '[{"name": "test_tool"}]', @@ -993,7 +1122,10 @@ class TestMCPToolManageService: } # Act: Execute the method under test - MCPToolManageService.update_mcp_provider( + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.update_provider( tenant_id=tenant.id, provider_id=mcp_provider.id, name="Updated MCP Provider", @@ -1015,7 +1147,10 @@ class TestMCPToolManageService: assert mcp_provider.updated_at is not None # Verify reconnection was called - mock_reconnect.assert_called_once_with("https://new-example.com/mcp", mcp_provider.id, tenant.id) + mock_reconnect.assert_called_once_with( + server_url="https://new-example.com/mcp", + provider=mcp_provider, + ) def test_update_mcp_provider_duplicate_name(self, db_session_with_containers, mock_external_service_dependencies): """ @@ -1048,8 +1183,11 @@ class TestMCPToolManageService: db.session.commit() # Act & Assert: Verify proper error handling for duplicate name + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="MCP tool First Provider already exists"): - MCPToolManageService.update_mcp_provider( + service.update_provider( tenant_id=tenant.id, provider_id=provider2.id, name="First Provider", # Duplicate name @@ -1094,19 +1232,22 @@ class TestMCPToolManageService: # Mock the provider controller and encryption with ( - patch("services.tools.mcp_tools_manage_service.MCPToolProviderController") as mock_controller, - patch("services.tools.mcp_tools_manage_service.ProviderConfigEncrypter") as mock_encrypter, + patch("core.tools.mcp_tool.provider.MCPToolProviderController") as mock_controller, + patch("core.tools.utils.encryption.ProviderConfigEncrypter") as mock_encrypter, ): # Setup mocks - mock_controller_instance = mock_controller._from_db.return_value + mock_controller_instance = mock_controller.from_db.return_value mock_controller_instance.get_credentials_schema.return_value = [] mock_encrypter_instance = mock_encrypter.return_value mock_encrypter_instance.encrypt.return_value = {"new_key": "encrypted_value"} # Act: Execute the method under test - MCPToolManageService.update_mcp_provider_credentials( - mcp_provider=mcp_provider, credentials={"new_key": "new_value"}, authed=True + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.update_provider_credentials( + provider=mcp_provider, credentials={"new_key": "new_value"}, authed=True ) # Assert: Verify the expected outcomes @@ -1152,19 +1293,22 @@ class TestMCPToolManageService: # Mock the provider controller and encryption with ( - patch("services.tools.mcp_tools_manage_service.MCPToolProviderController") as mock_controller, - patch("services.tools.mcp_tools_manage_service.ProviderConfigEncrypter") as mock_encrypter, + patch("core.tools.mcp_tool.provider.MCPToolProviderController") as mock_controller, + patch("core.tools.utils.encryption.ProviderConfigEncrypter") as mock_encrypter, ): # Setup mocks - mock_controller_instance = mock_controller._from_db.return_value + mock_controller_instance = mock_controller.from_db.return_value mock_controller_instance.get_credentials_schema.return_value = [] mock_encrypter_instance = mock_encrypter.return_value mock_encrypter_instance.encrypt.return_value = {"new_key": "encrypted_value"} # Act: Execute the method under test - MCPToolManageService.update_mcp_provider_credentials( - mcp_provider=mcp_provider, credentials={"new_key": "new_value"}, authed=False + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + service.update_provider_credentials( + provider=mcp_provider, credentials={"new_key": "new_value"}, authed=False ) # Assert: Verify the expected outcomes @@ -1199,14 +1343,18 @@ class TestMCPToolManageService: type("MockTool", (), {"model_dump": lambda self: {"name": "test_tool_2", "description": "Test tool 2"}})(), ] - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: # Setup mock client mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.return_value = mock_tools # Act: Execute the method under test - result = MCPToolManageService._re_connect_mcp_provider( - "https://example.com/mcp", mcp_provider.id, tenant.id + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service._reconnect_provider( + server_url="https://example.com/mcp", + provider=mcp_provider, ) # Assert: Verify the expected outcomes @@ -1224,16 +1372,8 @@ class TestMCPToolManageService: assert tools_data[1]["name"] == "test_tool_2" # Verify mock interactions - mock_mcp_client.assert_called_once_with( - "https://example.com/mcp", - mcp_provider.id, - tenant.id, - authed=False, - for_list=True, - headers={}, - timeout=30.0, - sse_read_timeout=300.0, - ) + provider_entity = mcp_provider.to_entity() + mock_mcp_client.assert_called_once() def test_re_connect_mcp_provider_auth_error(self, db_session_with_containers, mock_external_service_dependencies): """ @@ -1256,15 +1396,19 @@ class TestMCPToolManageService: ) # Mock MCPClient to raise authentication error - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: from core.mcp.error import MCPAuthError mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.side_effect = MCPAuthError("Authentication required") # Act: Execute the method under test - result = MCPToolManageService._re_connect_mcp_provider( - "https://example.com/mcp", mcp_provider.id, tenant.id + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) + result = service._reconnect_provider( + server_url="https://example.com/mcp", + provider=mcp_provider, ) # Assert: Verify the expected outcomes @@ -1295,12 +1439,18 @@ class TestMCPToolManageService: ) # Mock MCPClient to raise connection error - with patch("services.tools.mcp_tools_manage_service.MCPClient") as mock_mcp_client: + with patch("services.tools.mcp_tools_manage_service.MCPClientWithAuthRetry") as mock_mcp_client: from core.mcp.error import MCPError mock_client_instance = mock_mcp_client.return_value.__enter__.return_value mock_client_instance.list_tools.side_effect = MCPError("Connection failed") # Act & Assert: Verify proper error handling + from extensions.ext_database import db + + service = MCPToolManageService(db.session()) with pytest.raises(ValueError, match="Failed to re-connect MCP server: Connection failed"): - MCPToolManageService._re_connect_mcp_provider("https://example.com/mcp", mcp_provider.id, tenant.id) + service._reconnect_provider( + server_url="https://example.com/mcp", + provider=mcp_provider, + )