package keyword_test import ( "testing" "github.com/stretchr/testify/assert" "github.com/stretchr/testify/require" "github.com/yaoapp/yao/agent/context" "github.com/yaoapp/yao/agent/search/nlp/keyword" searchTypes "github.com/yaoapp/yao/agent/search/types" "github.com/yaoapp/yao/agent/testutils" oauthTypes "github.com/yaoapp/yao/openapi/oauth/types" ) func TestMCPProviderWithAssistantConfig(t *testing.T) { // Skip if running short tests if testing.Short() { t.Skip("Skipping integration test") } // Initialize test environment testutils.Prepare(t) defer testutils.Clean(t) // Create test context ctx := newMCPTestContext(t) // Create extractor with MCP mode extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{ MaxKeywords: 5, Language: "auto", }) // Test extraction content := "Machine learning and deep learning are subfields of artificial intelligence" keywords, err := extractor.Extract(ctx, content, nil) require.NoError(t, err) assert.NotEmpty(t, keywords, "MCP should return keywords") assert.LessOrEqual(t, len(keywords), 5, "Should respect max_keywords") t.Logf("Extracted keywords via MCP: %v", keywords) } func TestMCPProviderWithCustomOptions(t *testing.T) { // Skip if running short tests if testing.Short() { t.Skip("Skipping integration test") } // Initialize test environment testutils.Prepare(t) defer testutils.Clean(t) // Create test context ctx := newMCPTestContext(t) // Create extractor with MCP mode extractor := keyword.NewExtractor("mcp:search.extract_keywords", &searchTypes.KeywordConfig{ MaxKeywords: 10, }) // Test with runtime options override content := "Python programming language for data science and web development" keywords, err := extractor.Extract(ctx, content, &searchTypes.KeywordOptions{ MaxKeywords: 3, // Override to 3 }) require.NoError(t, err) assert.NotEmpty(t, keywords) assert.LessOrEqual(t, len(keywords), 3, "Should respect runtime max_keywords override") t.Logf("Extracted keywords via MCP (max 3): %v", keywords) } func TestMCPProviderInvalidFormat(t *testing.T) { // Test invalid MCP format fallback to system agent (requires context) extractor := keyword.NewExtractor("mcp:invalid", nil) // Should fallback to system agent which requires context _, err := extractor.Extract(nil, "test content for keyword extraction", nil) assert.Error(t, err) assert.Contains(t, err.Error(), "context is required") } func TestMCPProviderServerNotFound(t *testing.T) { // Skip if running short tests if testing.Short() { t.Skip("Skipping integration test") } // Initialize test environment testutils.Prepare(t) defer testutils.Clean(t) // Create test context ctx := newMCPTestContext(t) // Create extractor with non-existent MCP server extractor := keyword.NewExtractor("mcp:nonexistent.extract_keywords", &searchTypes.KeywordConfig{}) _, err := extractor.Extract(ctx, "test content", nil) assert.Error(t, err, "Should error for non-existent MCP server") assert.Contains(t, err.Error(), "not found") } func TestMCPProviderToolNotFound(t *testing.T) { // Skip if running short tests if testing.Short() { t.Skip("Skipping integration test") } // Initialize test environment testutils.Prepare(t) defer testutils.Clean(t) // Create test context ctx := newMCPTestContext(t) // Create extractor with non-existent tool extractor := keyword.NewExtractor("mcp:search.nonexistent_tool", &searchTypes.KeywordConfig{}) _, err := extractor.Extract(ctx, "test content", nil) assert.Error(t, err, "Should error for non-existent MCP tool") } func TestMCPProviderEmptyContent(t *testing.T) { // Skip if running short tests if testing.Short() { t.Skip("Skipping integration test") } // Initialize test environment testutils.Prepare(t) defer testutils.Clean(t) // Create test context ctx := newMCPTestContext(t) // Create extractor with MCP mode extractor := keyword.NewExtractor("mcp:search.extract_keywords", nil) // Test with empty content - MCP tool should return error _, err := extractor.Extract(ctx, "", nil) assert.Error(t, err, "Should error for empty content") } // newMCPTestContext creates a test context for MCP tests func newMCPTestContext(t *testing.T) *context.Context { t.Helper() authorized := &oauthTypes.AuthorizedInfo{ UserID: "test-user", } chatID := "test-chat-mcp-keyword" ctx := context.New(t.Context(), authorized, chatID) return ctx }